fix version guards

pytorch · Sep 26, 2024 · 46aeb9b · 46aeb9b
1 parent 5acff84
commit 46aeb9b
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/test/sparsity/test_sparse_api.py b/test/sparsity/test_sparse_api.py
@@ -13,7 +13,7 @@
 )
 
 from torchao.sparsity import apply_fake_sparsity, semi_sparse_weight, sparsify_
-from torchao.utils import TORCH_VERSION_AFTER_2_5, TORCH_VERSION_AT_LEAST_2_3
+from torchao.utils import TORCH_VERSION_AFTER_2_5, TORCH_VERSION_AT_LEAST_2_3, TORCH_VERSION_AT_LEAST_2_5, TORCH_VERSION_AT_LEAST_2_4
 
 
 logging.basicConfig(
@@ -48,7 +48,7 @@ def test_sparse(self):
 
 class TestQuantSemiSparse(common_utils.TestCase):
 
-    @unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_3, "pytorch 2.3+ feature")
+    @unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_4, "pytorch 2.4+ feature")
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
     @common_utils.parametrize("compile", [True, False])
     def test_quant_semi_sparse(self, compile):
@@ -79,6 +79,7 @@ def test_quant_semi_sparse(self, compile):
 
         torch.testing.assert_close(dense_result, sparse_result, rtol=1e-2, atol=1e-2)
 
+    @unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_5, "pytorch 2.5+ feature")
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
     @common_utils.parametrize("compile", [True, False])
     def test_sparse_marlin(self, compile):
@@ -110,7 +111,7 @@ def test_sparse_marlin(self, compile):
 
 
 class TestBlockSparseWeight(common_utils.TestCase):
-    @unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_3, "pytorch 2.3+ feature")
+    @unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_4, "pytorch 2.4+ feature due to need for custom op support")
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
     @common_utils.parametrize("compile", [True, False])
     def test_sparse(self, compile):