docker

huggingface · Feb 7, 2024 · 23de08d · 23de08d
1 parent 2b3dfb1
commit 23de08d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/models/llama/test_inference_llama.py b/tests/models/llama/test_inference_llama.py
@@ -26,7 +26,8 @@ def test_foo(self):
         model.to(device)
 
         # Generate
-        generate_ids = model.generate(inputs.input_ids, max_length=30)
+        with torch.no_grad():
+            generate_ids = model.generate(inputs.input_ids, max_length=30)
         output = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
 
         expected_output = "Hey, are you conscious? Can you talk to me?\nI'm here, but I'm not sure I'm conscious."