no_grad

thohag · Mar 15, 2023 · c671211 · c671211
1 parent 4dbebed
commit c671211
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/script.py b/script.py
@@ -37,14 +37,14 @@ def _call(self, prompt, stop=None):
             top_p=0.95,
             repetition_penalty=1.15,
         )
-
-        generation_output = model.generate(
-            input_ids=input_ids,
-            generation_config=generation_config,
-            return_dict_in_generate=True,
-            output_scores=True,
-            max_new_tokens=128,
-        )
+        with torch.no_grad():
+            generation_output = model.generate(
+                input_ids=input_ids,
+                generation_config=generation_config,
+                return_dict_in_generate=True,
+                output_scores=True,
+                max_new_tokens=128,
+            )
         response = ""
         for s in generation_output.sequences:
             response += tokenizer.decode(s)