Update i6_models/primitives/mixup.py

Co-authored-by: Albert Zeyer <[email protected]>
rwth-i6 · May 23, 2024 · 9c15fb5 · 9c15fb5
1 parent 74a1933
commit 9c15fb5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/i6_models/primitives/mixup.py b/i6_models/primitives/mixup.py
@@ -77,8 +77,8 @@ def get_random(self, b_dim: int, t_dim: int, max_num_mixup: int, n_mask: torch.t
                 start_indicies, n_mask
             )  # [B, M'] (M' denotes sum of num_mixup over the batch)
 
-            idx = torch.arange(t_dim)
-            idx = torch.unsqueeze(idx, dim=-1) + start_indicies_flat  # [T, M']
+            idx = torch.arange(t_dim)  # [T]
+            idx = idx[:, None] + start_indicies_flat[None, :]  # [T, M']
             mixup_values = self.cache[idx]  # [T, M', F]
             return mixup_values