allow extract backbone weights on non-GPU machines

gau-nernst · Jan 6, 2022 · 520940c · 520940c
1 parent 0b7a031
commit 520940c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/extras.py b/extras.py
@@ -108,9 +108,9 @@ def forward(self, batch, target):
 
 
 def extract_backbone_weights(lightning_ckpt_path, save_path):
-    ckpt = torch.load(lightning_ckpt_path)
+    ckpt = torch.load(lightning_ckpt_path, map_location='cpu')
     state_dict = ckpt["state_dict"]
-    backbone_weights = {k[len("backbone."):]: v.cpu() for k, v in state_dict.items() if k.startswith("backbone.")}
+    backbone_weights = {k[len("backbone."):]: v for k, v in state_dict.items() if k.startswith("backbone.")}
     torch.save(backbone_weights, save_path)