Update modeling_quiet.py

modeling_quiet.py

@@ -1426,15 +1426,12 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         return logits
     
     
+    @torch.no_grad()
     def generate(self, *args, **kwargs):
-        # Save the original input_ids and attention_mask
-        original_input_ids = kwargs.pop("input_ids", None)
-        original_attention_mask = kwargs.pop("attention_mask", None)
-
         # Call the infer method to get the logits
         logits = self.infer(
-            input_ids=original_input_ids,
-            attention_mask=original_attention_mask,
+            input_ids=kwargs.pop("input_ids", None),
+            attention_mask=kwargs.pop("attention_mask", None),
             position_ids=kwargs.pop("position_ids", None),
             past_key_values=kwargs.pop("past_key_values", None),
             inputs_embeds=kwargs.pop("inputs_embeds", None),