Update checkpoint.py (#722)

8eae7779 · nicheng0019 · GitHub · 3392a4ea · 8eae7779
Unverified Commit 8eae7779 authored Dec 18, 2020 by nicheng0019 Committed by GitHub Dec 18, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

mmcv/runner/checkpoint.py mmcv/runner/checkpoint.py +1 -1

No files found.
--- a/mmcv/runner/checkpoint.py
+++ b/mmcv/runner/checkpoint.py
@@ -256,7 +256,7 @@ def load_checkpoint(model,
        state_dict = checkpoint
    # strip prefix of state_dict
    if list(state_dict.keys())[0].startswith('module.'):
-        state_dict = {k[7:]: v for k, v in checkpoint['state_dict'].items()}
+        state_dict = {k[7:]: v for k, v in state_dict.items()}
    # load state_dict
    load_state_dict(model, state_dict, strict, logger)
    return checkpoint