From 08fedf4040ec7c460505222620ff35f932ae176e Mon Sep 17 00:00:00 2001 From: inkcherry Date: Fri, 30 Aug 2024 05:00:34 +0000 Subject: [PATCH] fix merge --- megatron/model/transformer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py index cc8581792a..617881e3a4 100644 --- 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)