remove elif branch

open-mmlab · Dec 12, 2022 · b267a16 · b267a16
1 parent 884fcad
commit b267a16
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/mmengine/model/weight_init.py b/mmengine/model/weight_init.py
@@ -487,14 +487,15 @@ def __call__(self, module):
                 module,
                 self.checkpoint,
                 map_location=self.map_location,
-                strict=False)
+                strict=False,
+                logger='current')
         else:
             print_log(
                 f'load {self.prefix} in model from: {self.checkpoint}',
                 logger='current')
             state_dict = _load_checkpoint_with_prefix(
                 self.prefix, self.checkpoint, map_location=self.map_location)
-            load_state_dict(module, state_dict, strict=False)
+            load_state_dict(module, state_dict, strict=False, logger='current')
 
         if hasattr(module, '_params_init_info'):
             update_init_info(module, init_info=self._get_init_info())

diff --git a/mmengine/runner/checkpoint.py b/mmengine/runner/checkpoint.py
@@ -107,10 +107,8 @@ def load(module, prefix=''):
         err_msg = '\n'.join(err_msg)
         if strict:
             raise RuntimeError(err_msg)
-        elif logger is not None:
-            logger.warning(err_msg)
         else:
-            print_log(err_msg, logger='current', level=logging.WARNING)
+            print_log(err_msg, logger=logger, level=logging.WARNING)
 
 
 def get_torchvision_models():