Forgot the X

allenai · Oct 26, 2023 · c8c6c68 · c8c6c68
1 parent b67e92c
commit c8c6c68
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/olmo/model.py b/olmo/model.py
@@ -972,7 +972,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: