[Fix] Fix optim_wrapper unittest for pytorch <= 1.10.0 (#975)

open-mmlab · Mar 2, 2023 · b1b1f53 · b1b1f53
1 parent 2ed8e34
commit b1b1f53
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/tests/test_optim/test_optimizer/test_optimizer_wrapper.py b/tests/test_optim/test_optimizer/test_optimizer_wrapper.py
@@ -17,6 +17,8 @@
 from mmengine.optim import AmpOptimWrapper, ApexOptimWrapper, OptimWrapper
 from mmengine.testing import assert_allclose
 from mmengine.testing._internal import MultiProcessTestCase
+from mmengine.utils.dl_utils import TORCH_VERSION
+from mmengine.utils.version_utils import digit_version
 
 is_apex_available = False
 try:
@@ -438,6 +440,10 @@ def test_init(self):
         not torch.cuda.is_available(),
         reason='`torch.cuda.amp` is only available when pytorch-gpu installed')
     def test_step(self, dtype):
+        if dtype is not None and (digit_version(TORCH_VERSION) <
+                                  digit_version('1.10.0')):
+            raise unittest.SkipTest('Require PyTorch version >= 1.10.0 to '
+                                    'support `dtype` argument in autocast')
         if dtype == 'bfloat16' and not bf16_supported():
             raise unittest.SkipTest('bfloat16 not supported by device')
         optimizer = MagicMock(spec=Optimizer)
@@ -454,6 +460,10 @@ def test_step(self, dtype):
         not torch.cuda.is_available(),
         reason='`torch.cuda.amp` is only available when pytorch-gpu installed')
     def test_backward(self, dtype):
+        if dtype is not None and (digit_version(TORCH_VERSION) <
+                                  digit_version('1.10.0')):
+            raise unittest.SkipTest('Require PyTorch version >= 1.10.0 to '
+                                    'support `dtype` argument in autocast')
         if dtype == 'bfloat16' and not bf16_supported():
             raise unittest.SkipTest('bfloat16 not supported by device')
         amp_optim_wrapper = AmpOptimWrapper(
@@ -512,6 +522,10 @@ def test_load_state_dict(self):
         not torch.cuda.is_available(),
         reason='`torch.cuda.amp` is only available when pytorch-gpu installed')
     def test_optim_context(self, dtype, target_dtype):
+        if dtype is not None and (digit_version(TORCH_VERSION) <
+                                  digit_version('1.10.0')):
+            raise unittest.SkipTest('Require PyTorch version >= 1.10.0 to '
+                                    'support `dtype` argument in autocast')
         if dtype == 'bfloat16' and not bf16_supported():
             raise unittest.SkipTest('bfloat16 not supported by device')
         amp_optim_wrapper = AmpOptimWrapper(