Update for the revision

events.out.tfevents.1626448850.t1v-n-278acf21-w-0.590260.3.v2

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1fbe385b41508eae766e3ae9763a6bf8a20b0dad2a36c5058b526b6884a8433a
-size 662195
+oid sha256:0f8f2848f118433d3ae3412ed5ed0df7242cdf899879357f922313aeaf0b7b5d
+size 809333

flax_model.msgpack

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:dd33994b480ef0a93c7821a12df82c34656dc30539b623c1fb2050b1ba03be19
-size 190539834

src/run_wav2vec2_pretrain_flax.py

@@ -160,7 +160,6 @@ class FlaxDataCollatorForWav2Vec2Pretraining:
     """
     Data collator that will dynamically pad the inputs received and prepare masked indices
     for self-supervised pretraining.
-
     Args:
         model (:class:`~transformers.FlaxWav2Vec2ForPreTraining`):
             The Wav2Vec2 model used for pretraining. The data collator needs to have access
@@ -203,6 +202,7 @@ class FlaxDataCollatorForWav2Vec2Pretraining:
 
         batch_size = batch["input_values"].shape[0]
 
+        attention_mask = None
         if batch["attention_mask"] is not None:
             output_lengths = self.model._get_feat_extract_output_lengths(batch["attention_mask"].sum(-1))
             attention_mask = np.zeros((batch_size, mask_indices_seq_length), dtype=np.int8)
@@ -225,9 +225,11 @@ class FlaxDataCollatorForWav2Vec2Pretraining:
         batch["sampled_negative_indices"] = _sample_negative_indices(
             (batch["mask_time_indices"].shape + (self.model.config.proj_codevector_dim,)),
             self.model.config.num_negatives,
+            attention_mask=attention_mask,
         )
 
         return batch
+        
 
 def configure_logger(model_args: ModelArguments, training_args: TrainingArguments):
     logging.basicConfig(