Update modeling.py

modeling.py

@@ -185,6 +185,9 @@ class RecombinationTransformerForCausalLM(PreTrainedModel):
         if past:
             input_ids = input_ids[:, -1].unsqueeze(-1)
         
+        if attention_mask is None:
+            attention_mask = torch.ones(input_ids.shape, device=input_ids.device)
+        
         return {"input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past}
         
     def generate(self, input_ids, attention_mask=None, max_length=512, min_length=None, num_return_sequences=1):