Merge pull request #24 from xingchensong/xcsong-fix-mask

[model] fix mask
xingchensong · Dec 26, 2024 · 8169702 · 8169702
2 parents c0954be + 3f67a78
commit 8169702
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/s3tokenizer/model.py b/s3tokenizer/model.py
@@ -178,10 +178,10 @@ def forward(self, x: Tensor, x_len: Tensor) -> Tuple[Tensor, Tensor]:
         x = F.gelu(self.conv2(x))
         x = x.permute(0, 2, 1)  # (B, T // 2, n_state)
         mask = make_non_pad_mask(x_len, T).unsqueeze(1)  # (B, 1, T)
-        mask = mask[:, :, (T + 1) % 2::2]  # (B, 1, T // 2)
+        mask = mask[:, :, (T + 2) % 2::2]  # (B, 1, T // 2)
         if self.stride == 2:
             _T = mask.size(-1)
-            mask = mask[:, :, (_T + 1) % 2::2]  # (B, 1, T // 4)
+            mask = mask[:, :, (_T + 2) % 2::2]  # (B, 1, T // 4)
         mask = mask_to_bias(mask, x.dtype)
 
         x = (x + self.positional_embedding[:x.shape[1], :]).to(x.dtype)

diff --git a/s3tokenizer/model_v2.py b/s3tokenizer/model_v2.py
@@ -303,7 +303,7 @@ def forward(self, x: torch.Tensor,
         x = x.permute(0, 2, 1)  # (B, T // 2, n_state)
         freqs_cis = self.freqs_cis.to(x.device)
         mask = make_non_pad_mask(x_len, T).unsqueeze(1)  # (B, 1, T)
-        mask = mask[:, :, (T + 1) % 2::2]  # (B, 1, T // 2)
+        mask = mask[:, :, (T + 2) % 2::2]  # (B, 1, T // 2)
         mask_pad = None
         if self.stride == 2:
             _T = mask.size(-1)