Merge branch 'Llama-DefaultLN' of https://github.com/allenai/LLM into…

… Llama-DefaultLN
allenai · Oct 25, 2023 · bb8e2f6 · bb8e2f6
2 parents 06558d4 + e9c92c8
commit bb8e2f6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/olmo/model.py b/olmo/model.py
@@ -948,7 +948,7 @@ def forward(
         ):
             # shape: (batch_size, seq_len, d_model)
             x, cache = self.__activation_checkpoint_fn(
-                block, attention_bias=attention_bias, layer_past=layer_past, use_cache=use_cache
+                block, x, attention_bias=attention_bias, layer_past=layer_past, use_cache=use_cache
             )
 
             if attn_key_values is not None: