[Feature]: Change map_location to cpu

2022-12-09 14:40:40 +08:00 · 2022-12-09 14:40:40 +08:00 · 976de2e518
parent 5244b66d5c
commit 976de2e518
1 changed files with 1 additions and 1 deletions
--- a/tools/model_converters/mmcls2timm.py
+++ b/tools/model_converters/mmcls2timm.py
@ -71,7 +71,7 @@ def main():
    parser.add_argument('dst', help='save path')
    args = parser.parse_args()

-    checkpoint = _load_checkpoint(args.src)
+    checkpoint = _load_checkpoint(args.src, map_location='cpu')
    if 'state_dict' in checkpoint:
        state_dict = checkpoint['state_dict']
    else: