docker

huggingface · Feb 5, 2024 · 371d815 · 371d815
1 parent e5ad8e3
commit 371d815
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/tests/models/llama/test_inference_llama.py b/tests/models/llama/test_inference_llama.py
@@ -1,5 +1,6 @@
 import os
 import unittest
+import torch
 
 from transformers import AutoModelForCausalLM
 
@@ -15,5 +16,5 @@ def test_foo(self):
         assert os.getenv("HF_HUB_READ_TOKEN", None) != None
         assert type(os.getenv("HF_HUB_READ_TOKEN", None)) == str
         assert len(os.getenv("HF_HUB_READ_TOKEN", None)) > 3
-        model = AutoModelForCausalLM.from_pretrained(ckpt, token=os.getenv("HF_HUB_READ_TOKEN", None))
+        model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16, token=os.getenv("HF_HUB_READ_TOKEN", None))
         model.to(device)
diff --git a/tests/models/llama/test_train_llama.py b/tests/models/llama/test_train_llama.py
@@ -1,5 +1,6 @@
 import os
 import unittest
+import torch
 
 from transformers import AutoModelForCausalLM
 
@@ -13,6 +14,6 @@ def test_foo(self):
         assert os.getenv("HF_HUB_READ_TOKEN", None) != None
         assert type(os.getenv("HF_HUB_READ_TOKEN", None)) == str
         assert len(os.getenv("HF_HUB_READ_TOKEN", None)) > 3
-        model = AutoModelForCausalLM.from_pretrained(ckpt, token=os.getenv("HF_HUB_READ_TOKEN", None))
+        model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16, token=os.getenv("HF_HUB_READ_TOKEN", None))
         model.train()
         model.to(device)