fix: add model_id to model test

huggingface · Jun 6, 2024 · 4f90ff3 · 4f90ff3
1 parent 60a0956
commit 4f90ff3
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/server/tests/models/test_model.py b/server/tests/models/test_model.py
@@ -17,7 +17,12 @@ def generate_token(self, batch):
     tokenizer = AutoTokenizer.from_pretrained("huggingface/llama-7b")
 
     model = TestModel(
-        torch.nn.Linear(1, 1), tokenizer, False, torch.float32, torch.device("cpu")
+        "test_model_id",
+        torch.nn.Linear(1, 1),
+        tokenizer,
+        False,
+        torch.float32,
+        torch.device("cpu"),
     )
     return model