Fix cuda-cpu load in convert

awi-response · Jan 9, 2025 · 21d295f · 21d295f
1 parent ff2999f
commit 21d295f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/darts/src/darts/legacy_training/util.py b/darts/src/darts/legacy_training/util.py
@@ -33,7 +33,7 @@ def convert_lightning_checkpoint(
     import torch
 
     logger.debug(f"Loading checkpoint from {lightning_checkpoint.resolve()}")
-    lckpt = torch.load(lightning_checkpoint, weights_only=False)
+    lckpt = torch.load(lightning_checkpoint, weights_only=False, map_location=torch.device("cpu"))
 
     now = datetime.now()
     formatted_date = now.strftime("%Y-%m-%d")