fix bug for flash attention (EleutherAI#910)

Jarvis-LLM · May 9, 2023 · dc05783 · dc05783
1 parent befd133
commit dc05783
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
@@ -542,7 +542,7 @@ def flash_attention(self, query_layer, key_layer, value_layer):
  else:
 
  # [sq, b, np, hn] -> [b * sq, 1, np, hn]
- query_layer.transpose(0, 1).reshape(
+ query_layer = query_layer.transpose(0, 1).reshape(
  output_size[0] * output_size[2], 1, output_size[1], -1
  )