fixing tests:...

neuralmagic · Dec 29, 2023 · 682fd9d · 682fd9d
1 parent 92c6017
commit 682fd9d
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/tests/sparseml/pytorch/conftest.py b/tests/sparseml/pytorch/conftest.py
@@ -11,13 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from sparseml.pytorch.image_classification.integration_helper_functions import (
-    ImageClassification,
-)
 
 
 def test_integration_helper_functions():
     yield
+    from sparseml.pytorch.image_classification.integration_helper_functions import (
+        ImageClassification,
+    )
+
     image_classification = ImageClassification()
     assert image_classification.create_model
     assert image_classification.create_dummy_input