Revert "Fix Llama initialization"

This reverts commit ddb4c02.
huggingface · Feb 14, 2024 · 18252c7 · 18252c7
1 parent ddb4c02
commit 18252c7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py
@@ -414,7 +414,7 @@ def post_mlp_forward(self, x):
 
 class GaudiLlamaDecoderLayer(LlamaDecoderLayer):
     def __init__(self, config: LlamaConfig, layer_idx: int):
-        super(LlamaDecoderLayer, self).__init__()
+        super().__init__(config, layer_idx)
         self.hidden_size = config.hidden_size
 
         self.self_attn = GaudiLlamaAttention(config=config, layer_idx=layer_idx)
@@ -666,7 +666,7 @@ def forward(
                     attn_softmax_bf16,
                     False,
                     use_flash_attention,
-                    flash_attention_recompute,
+                    True,
                 )
             else:
                 layer_outputs = decoder_layer(