refactor

xrsrke · Sep 1, 2023 · 55533cb · 55533cb
1 parent 9401783
commit 55533cb
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/tests/nn/tensor_parallel/conftest.py b/tests/nn/tensor_parallel/conftest.py
@@ -0,0 +1,9 @@
+import pytest
+from transformers import AutoModel
+
+MODEL_NAME = "bigscience/bloom-560m"
+
+
+@pytest.fixture(scope="session")
+def model():
+    return AutoModel.from_pretrained(MODEL_NAME)