Update modeling_molmo.py for compatibility

modeling_molmo.py

@@ -2273,6 +2273,11 @@ class MolmoForCausalLM(PreTrainedModel):
                 del model_kwargs["image_masks"]
                 del model_kwargs["image_input_idx"]
         cache_name, cache = super()._extract_past_from_model_output(outputs)
+        try:
+            cache_name, cache = super()._extract_past_from_model_output(outputs)
+        except AttributeError:
+            past_key_values = outputs.past_key_values if "past_key_values" in outputs else None
+            cache_name, cache = "past_key_values", past_key_values
         model_kwargs[cache_name] = cache
         model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
         return model_kwargs