Update modeling_internlm2.py

modeling_internlm2.py

@@ -1088,7 +1088,7 @@ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
     def prepare_inputs_for_generation(
         self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
     ):
-        if past_key_values is not None:
+        if past_key_values is not None and past_key_values[0][0] is not None:
             past_length = past_key_values[0][0].shape[2]
 
             # Some generation methods already pass only the last input ID