[Fix]: fix load_checkpoint (open-mmlab#332)

HAOCHENYE · Jun 23, 2022 · 12f7d3a · 12f7d3a
1 parent 2994195
commit 12f7d3a
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/mmengine/runner/runner.py b/mmengine/runner/runner.py
@@ -1806,13 +1806,8 @@ def resume(self,
                 Defaults to 'default'.
         """
         if map_location == 'default':
-            if torch.cuda.is_available():
-                device = get_device()
-                checkpoint = self.load_checkpoint(
-                    filename,
-                    map_location=lambda storage, loc: storage.to(device))
-            else:
-                checkpoint = self.load_checkpoint(filename)
+            device = get_device()
+            checkpoint = self.load_checkpoint(filename, map_location=device)
         else:
             checkpoint = self.load_checkpoint(
                 filename, map_location=map_location)