[bestrq] fix time stack position (#1869)

* [bestrq] fix time stack position fix time stack position * fix lint * fix mask tail of batch
wenet-e2e · May 30, 2023 · abcc9c5 · abcc9c5
1 parent 2505d18
commit abcc9c5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/wenet/ssl/bestrq/bestqr_model.py b/wenet/ssl/bestrq/bestqr_model.py
@@ -191,7 +191,7 @@ def forward(
         subsampling_masks = masked_masks.unfold(1,
                                                 size=self.stack_frames,
                                                 step=self.stride)
-        code_ids_mask, _ = torch.max(subsampling_masks, 2)
+        code_ids_mask, _ = torch.min(subsampling_masks, 2)
 
         # 2.0 stack fbank
         unmasked_xs = self._stack_features(input)
@@ -251,6 +251,7 @@ def _apply_mask_signal(
     def _stack_features(self, input: torch.Tensor) -> torch.Tensor:
 
         stack_input = input.unfold(1, size=self.stack_frames, step=self.stride)
+        stack_input = stack_input.transpose(-1, -2)
         b, n, f, d = stack_input.size()
         stack_input = stack_input.reshape(b, n, f * d)