Update i6_models/parts/best_rq/mask.py

Co-authored-by: michelwi <[email protected]>
rwth-i6 · Dec 4, 2024 · b50b66e · b50b66e
1 parent b34b1dc
commit b50b66e
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/i6_models/parts/best_rq/mask.py b/i6_models/parts/best_rq/mask.py
@@ -66,13 +66,8 @@ def forward(
                 min_len = seq_len - num_mask - 1
             mask_idc = np.random.choice(seq_len - min_len, num_mask, replace=False)
 
-            mask_idc = np.asarray(
-                [mask_idc[j] + offset for j in range(len(mask_idc)) for offset in range(self.mask_length)]
-            )
-            mask_idcs.append(mask_idc)
-
-        for i, mask_idc in enumerate(mask_idcs):
-            mask[i, mask_idc] = True
+            for j in mask_idc:
+                mask[i, j : j+self.mask_length] = True
 
         tensor[mask] = self.mask_emb.to(tensor.device)