add examples

wenet-e2e · Sep 27, 2023 · a392d05 · a392d05
1 parent dda5a69
commit a392d05
Showing 1 changed file with 17 additions and 2 deletions.
diff --git a/wenet/utils/mask.py b/wenet/utils/mask.py
@@ -322,7 +322,7 @@ def causal_or_lookahead_mask(
         [1, 1, 1, 0],
         [1, 1, 1, 1]]
         >>> causal_or_lookahead_mask(seq_mask.unsqueeze(1), 0, 2)
-        [[1, 0, 0, 0],
+        [[[1, 0, 0, 0],
          [1, 1, 0, 0],
          [0, 0, 0, 0],
          [0, 0, 0, 0]],
@@ -335,7 +335,22 @@ def causal_or_lookahead_mask(
         [[1, 0, 0, 0],
          [1, 1, 0, 0],
          [1, 1, 1, 0],
-         [0, 1, 1, 1]]
+         [0, 1, 1, 1]]]
+        >>> causal_or_lookahead_mask(seq_mask.unsqueeze(1), 1, 2)
+        [[[1, 1, 0, 0],
+         [1, 1, 0, 0],
+         [0, 0, 0, 0],
+         [0, 0, 0, 0]],
+
+        [[1, 1, 0, 0],
+         [1, 1, 1, 0],
+         [1, 1, 1, 0],
+         [0, 0, 0, 0]],
+
+        [[1, 1, 0, 0],
+         [1, 1, 1, 0],
+         [1, 1, 1, 1],
+         [0, 1, 1, 1]]]
     """
     _, _, T = mask.size()
     indices = torch.arange(T, device=mask.device)