

TITLE
[Wav2Vec2] Improve SpecAugment function by converting numpy based function to pytorch based function

COMMENTS
17

REACTIONS
+1: 2
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
# 🚀 Feature request

As can be seen here: https://github.com/huggingface/transformers/blob/11655fafdd42eb56ad94e09ecd84d4dc2d1041ae/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L47,

the function `_compute_mask_indices` (responsible for spec augment) of Wav2Vec2 is written in numpy which means that the function is not GPU compatible. 
The function could simply be rewritten in PyTorch, which should make training on GPU faster.

This "Good First Issue" is about converting `_compute_mask_indices` to PyTorch while keeping the same functionality. 

## Your contribution

I'm happy to guide the contributor through the PR!
