fix bug

open-mmlab · Jul 26, 2023 · 78c8e07 · 78c8e07
1 parent a802a3d
commit 78c8e07
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tools/train.py b/tools/train.py
@@ -98,7 +98,8 @@ def merge_args(cfg, args):
     if args.amp is True:
         from mmengine.optim import AmpOptimWrapper, OptimWrapper
         optim_wrapper = cfg.optim_wrapper.get('type', OptimWrapper)
-        assert optim_wrapper in (OptimWrapper, AmpOptimWrapper), \
+        assert optim_wrapper in (OptimWrapper, AmpOptimWrapper,
+                                 'OptimWrapper', 'AmpOptimWrapper'), \
             '`--amp` is not supported custom optimizer wrapper type ' \
             f'`{optim_wrapper}.'
         cfg.optim_wrapper.type = 'AmpOptimWrapper'