Merge branch 'main' of https://github.com/George0828Zhang/torch_cif

George0828Zhang · Jun 10, 2022 · 6970623 · 6970623
2 parents 331ccb1 + aaeb53c
commit 6970623
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/cif.py b/cif.py
@@ -173,10 +173,11 @@ def cif_function(
                 .scatter(
                     1,
                     feat_lengths.view(B, 1, 1).expand(-1, -1, C),
-                    beta / tail_weights.view(B, 1, 1).expand(-1, -1, C),
+                    beta / tail_weights.masked_fill(~extend_mask, beta).view(B, 1, 1).expand(-1, -1, C),
                 )
+                .detach()
             )
-            output[extend_mask] *= upscale[extend_mask]
+            output *= upscale
             feat_lengths += extend_mask.long()
             T = feat_lengths.max()
         output = output[:, :T, :]