fix pytest

tests/test_models/test_algorithms/test_algorithm.py

@@ -5,8 +5,9 @@
 
 import mmcv
 import numpy as np
+import pytest
 import torch
-from mmcv import Config, ConfigDict
+from mmcv import Config, ConfigDict, digit_version
 
 from mmrazor.models.builder import ALGORITHMS
 
@@ -95,6 +96,14 @@ def test_autoslim_pretrain():
         pruner=pruner_cfg,
         distiller=distiller_cfg)
 
+    # ``StructurePruner`` requires pytorch>=1.6.0 to
+    # auto-trace correctly
+    min_required_version = '1.6.0'
+    if digit_version(torch.__version__) < digit_version(min_required_version):
+        with pytest.raises(AssertionError):
+            model = ALGORITHMS.build(algorithm_cfg)
+        return
+
     imgs = torch.randn(16, 3, 224, 224)
     label = torch.randint(0, 1000, (16, ))
 
@@ -147,6 +156,14 @@ def test_autoslim_retrain():
         retraining=True,
         channel_cfg=channel_cfg)
 
+    # ``StructurePruner`` requires pytorch>=1.6.0 to
+    # auto-trace correctly
+    min_required_version = '1.6.0'
+    if digit_version(torch.__version__) < digit_version(min_required_version):
+        with pytest.raises(AssertionError):
+            model = ALGORITHMS.build(algorithm_cfg)
+        return
+
     imgs = torch.randn(16, 3, 224, 224)
     label = torch.randint(0, 1000, (16, ))
 

tests/test_models/test_pruner.py

@@ -3,7 +3,7 @@
 
 import pytest
 import torch
-from mmcv import ConfigDict
+from mmcv import ConfigDict, digit_version
 
 from mmrazor.models.builder import ARCHITECTURES, PRUNERS
 
@@ -35,6 +35,14 @@ def test_ratio_pruner():
         type='RatioPruner',
         ratios=[1 / 8, 2 / 8, 3 / 8, 4 / 8, 5 / 8, 6 / 8, 7 / 8, 1.0])
 
+    # ``StructurePruner`` requires pytorch>=1.6.0 to
+    # auto-trace correctly
+    min_required_version = '1.6.0'
+    if digit_version(torch.__version__) < digit_version(min_required_version):
+        with pytest.raises(AssertionError):
+            pruner = PRUNERS.build(pruner_cfg)
+        return
+
     _test_reset_bn_running_stats(architecture_cfg, pruner_cfg, False)
     with pytest.raises(AssertionError):
         _test_reset_bn_running_stats(architecture_cfg, pruner_cfg, True)