fix: add model_id to model test

2024-06-06 21:24:29 +00:00 · 2024-06-06 21:24:29 +00:00 · 9c45d34983
parent dc0f76553c
commit 9c45d34983
1 changed files with 6 additions and 1 deletions
--- a/server/tests/models/test_model.py
+++ b/server/tests/models/test_model.py
@ -17,7 +17,12 @@ def get_test_model():
    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