From 12f7d3a0d3470110b875bb1125b358e6c1d34563 Mon Sep 17 00:00:00 2001
From: Jiazhen Wang <47851024+teamwong111@users.noreply.github.com>
Date: Thu, 23 Jun 2022 16:53:53 +0800
Subject: [PATCH] [Fix]: fix load_checkpoint (#332)
---
mmengine/runner/runner.py | 9 ++-------
1 file changed, 2 insertions(+), 7 deletions(-)
diff --git a/mmengine/runner/runner.py b/mmengine/runner/runner.py
index fec759be..f754e88c 100644
--- a/mmengine/runner/runner.py
+++ b/mmengine/runner/runner.py
@@ -1806,13 +1806,8 @@ class Runner:
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)
--
GitLab