docker

huggingface · Feb 7, 2024 · e6eb5a6 · e6eb5a6
1 parent 23de08d
commit e6eb5a6
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/models/llama/test_inference_llama.py b/tests/models/llama/test_inference_llama.py
@@ -25,6 +25,9 @@ def test_foo(self):
         model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16, token=hf_token)
         model.to(device)
 
+        # To make generation's sample deterministic
+        torch.manual_seed(0)
+
         # Generate
         with torch.no_grad():
             generate_ids = model.generate(inputs.input_ids, max_length=30)