Update modeling_molmo.py

#41
Files changed (1) hide show
  1. modeling_molmo.py +1 -2
modeling_molmo.py CHANGED
@@ -2272,8 +2272,7 @@ class MolmoForCausalLM(PreTrainedModel):
2272
  del model_kwargs["images"]
2273
  del model_kwargs["image_masks"]
2274
  del model_kwargs["image_input_idx"]
2275
- cache_name, cache = super()._extract_past_from_model_output(outputs)
2276
- model_kwargs[cache_name] = cache
2277
  model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
2278
  return model_kwargs
2279
 
 
2272
  del model_kwargs["images"]
2273
  del model_kwargs["image_masks"]
2274
  del model_kwargs["image_input_idx"]
2275
+ model_kwargs["past_key_values"] = outputs.past_key_values
 
2276
  model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
2277
  return model_kwargs
2278