restore bestrq

wenet-e2e · Sep 20, 2024 · 2dc8baa · 2dc8baa
1 parent 42798bc
commit 2dc8baa
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/wenet/ssl/bestrq/bestrq_model.py b/wenet/ssl/bestrq/bestrq_model.py
@@ -279,7 +279,6 @@ def _stack_features(self, input: torch.Tensor,
     def _compute_loss(self, input: torch.Tensor, target: torch.Tensor,
                       mask: torch.Tensor) -> torch.Tensor:
         logits = input.transpose(1, 2).contiguous().view(-1, input.size(-1))
-        mask = mask.unsqueeze(2).repeat(1, 1, self.num_codebooks)
         loss = torch.nn.functional.cross_entropy(
             logits,
             target.contiguous().view(-1),