Forgot the X

allenai · Oct 25, 2023 · e9c92c8 · e9c92c8
1 parent a9d5d6a
commit e9c92c8
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: