update assert

neuralmagic · Apr 9, 2024 · 6dab7dd · 6dab7dd
1 parent d7c3ae0
commit 6dab7dd
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/tests/sparseml/modifiers/quantization/test_base.py b/tests/sparseml/modifiers/quantization/test_base.py
@@ -62,8 +62,8 @@ def test_end_epochs(self):
             freeze_bn_stats_epoch=freeze_bn_epoch,
         )
 
-        assert obj_modifier.calculate_disable_observer_epoch() == -1
-        assert obj_modifier.calculate_freeze_bn_stats_epoch() == -1
+        self.assertEqual(obj_modifier.calculate_disable_observer_epoch(), -1)
+        self.assertEqual(obj_modifier.calculate_freeze_bn_stats_epoch(), -1)
 
         for epoch in range(3):
             event = Event(steps_per_epoch=1, global_step=epoch)

diff --git a/tests/sparseml/pytorch/modifiers/obcq/test_pytorch.py b/tests/sparseml/pytorch/modifiers/obcq/test_pytorch.py
@@ -174,7 +174,7 @@ def test_set_quant_in_sparsegpt(self):
         testing_harness = LifecyleTestingHarness(model=LinearNet())
         modifier.on_initialize_structure(testing_harness.get_state())
         assert modifier.quantize
-        assert isinstance(modifier.quantization_modifier_, QuantizationModifier)
+        self.assertIsInstance(modifier.quantization_modifier_, QuantizationModifier)
 
         dict_scheme = dict(modifier.quantization_modifier_.scheme)
         self.assertEqual(