missing import

neuralmagic · May 20, 2024 · 3e0669a · 3e0669a
1 parent 0539ef7
commit 3e0669a
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/sparseml/transformers/obcq/test_obcq_completion.py b/tests/sparseml/transformers/obcq/test_obcq_completion.py
@@ -148,6 +148,7 @@ class TestOBCQCompletionGPU(TestOBCQCompletion):
 
     def setUp(self):
         from sparseml.transformers import SparseAutoModelForCausalLM
+        import torch
 
         self.model_name = None
         self.output = "./oneshot_output"