docker

huggingface · Feb 2, 2024 · 35130d8 · 35130d8
1 parent fb478fe
commit 35130d8
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/tests/models/llama/test_inference_llama.py b/tests/models/llama/test_inference_llama.py
@@ -1,7 +1,14 @@
 import unittest
 
+from transformers import AutoModelForCausalLM
+
+device = "cuda"
+
 
 class LLaMaInferenceTest(unittest.TestCase):
 
     def test_foo(self):
-        assert 1 == 1
+        ckpt = "meta-llama/Llama-2-7b"
+        model = AutoModelForCausalLM.from_pretrained(ckpt)
+        model.train()
+        model.to(device)
diff --git a/tests/models/llama/test_train_llama.py b/tests/models/llama/test_train_llama.py
@@ -1,7 +1,14 @@
 import unittest
 
+from transformers import AutoModelForCausalLM
+
+device = "cuda"
+
 
 class LLaMaTrainingTest(unittest.TestCase):
 
     def test_foo(self):
-        assert 1 == 2
+        ckpt = "meta-llama/Llama-2-7b"
+        model = AutoModelForCausalLM.from_pretrained(ckpt)
+        model.train()
+        model.to(device)