fix merge

deepspeedai · Aug 30, 2024 · 08fedf4 · 08fedf4
1 parent 191fbf7
commit 08fedf4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
@@ -875,7 +875,7 @@ def forward(self, hidden_states, attention_mask,
                 if not self.use_flash_attn_triton:
                     query_layer, key_layer, value_layer = [rearrange(x, 's b ... -> b s ...').contiguous()
                             for x in (query_layer, key_layer, value_layer)]
-                batch_dim_idx = 0
+                    batch_dim_idx = 0
 
                 context_layer = self.dist_attn(query_layer, key_layer, value_layer, batch_dim_idx)