fix mask shape

pytorch · Aug 23, 2024 · b0454c3 · b0454c3
1 parent 2869d98
commit b0454c3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchtune/modules/transformer.py b/torchtune/modules/transformer.py
@@ -442,7 +442,7 @@ def forward(
                 )
             # shape: [1, input_pos_len, m_s]
             # in most cases input_pos_len should be 1
-            mask = self.causal_mask[None, input_pos]
+            mask = self.causal_mask[None, input_pos, :input_pos.max().item() + 1]
 
         hidden = []
         for i, layer in enumerate(self.layers):