Update modeling_molmo.py

modeling_molmo.py

@@ -2272,8 +2272,7 @@ class MolmoForCausalLM(PreTrainedModel):
                 del model_kwargs["images"]
                 del model_kwargs["image_masks"]
                 del model_kwargs["image_input_idx"]
-        cache_name, cache = super()._extract_past_from_model_output(outputs)
-        model_kwargs[cache_name] = cache
+        model_kwargs["past_key_values"] = outputs.past_key_values
         model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
         return model_kwargs