fix UT

nijkah · Sep 27, 2022 · dd64538 · dd64538
1 parent a408f69
commit dd64538
Showing 1 changed file with 4 additions and 11 deletions.
diff --git a/tests/test_optim/test_optimizer/test_optimizer.py b/tests/test_optim/test_optimizer/test_optimizer.py
@@ -8,6 +8,7 @@
 import torch
 import torch.nn as nn
 
+from mmengine.model import MMDistributedDataParallel
 from mmengine.optim import (OPTIM_WRAPPER_CONSTRUCTORS, OPTIMIZERS,
                             DefaultOptimWrapperConstructor, OptimWrapper,
                             build_optim_wrapper)
@@ -759,6 +760,7 @@ def _check_default_optimizer(self, optimizer, model, prefix=''):
     def test_build_zero_redundancy_optimizer(self):
         self._init_dist_env(self.rank, self.world_size)
         model = ExampleModel()
+        model = MMDistributedDataParallel(module=model)
         self.base_lr = 0.01
         self.momentum = 0.0001
         self.base_wd = 0.9
@@ -773,17 +775,8 @@ def test_build_zero_redundancy_optimizer(self):
                 momentum=self.momentum))
         optim_wrapper = build_optim_wrapper(model, optim_wrapper_cfg)
         self.assertIsInstance(optim_wrapper.optimizer, ZeroRedundancyOptimizer)
-        self._check_default_optimizer(optim_wrapper.optimizer, model)
-
-        # test build optimizer without ``optimizer_type``
-        with self.assertRaises(TypeError):
-            optim_wrapper_cfg = dict(
-                optimizer=dict(
-                    type='ZeroRedundancyOptimizer',
-                    lr=self.base_lr,
-                    weight_decay=self.base_wd,
-                    momentum=self.momentum))
-            optim_wrapper = build_optim_wrapper(model, optim_wrapper_cfg)
+        self._check_default_optimizer(
+            optim_wrapper.optimizer, model, prefix='module.')
 
     def _init_dist_env(self, rank, world_size):
         """Initialize the distributed environment."""