Update modeling_quiet.py

modeling_quiet.py

@@ -2174,6 +2174,8 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 
         if attention_mask is None:
             attention_mask = torch.ones_like(input_ids)
+        else:
+            attention_mask = attention_mask[:, -input_ids.shape[1]:]  # Adjust the attention mask size
 
         if self.use_start_thought_token:
             start_thought_token_id = self.tokenizer.convert_tokens_to_ids("<|startthought|>")