load ckpt to cpu (#428)

* load ckpt to cpu * precommit --------- Co-authored-by: Benjamin Morris <[email protected]>
AllenCellModeling · Oct 1, 2024 · 49b0bc0 · 49b0bc0
1 parent 48e7cb9
commit 49b0bc0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cyto_dl/utils/checkpoint.py b/cyto_dl/utils/checkpoint.py
@@ -7,7 +7,7 @@ def load_checkpoint(model, load_params):
             "ckpt_path"
         ), "ckpt_path must be provided to with argument weights_only=True"
         # load model from state dict to get around trainer.max_epochs limit, useful for resuming model training from existing weights
-        state_dict = torch.load(load_params["ckpt_path"])["state_dict"]
+        state_dict = torch.load(load_params["ckpt_path"], map_location="cpu")["state_dict"]
         model.load_state_dict(state_dict, strict=load_params.get("strict", True))
         # set ckpt_path to None to avoid loading checkpoint again with model.fit/model.test
         load_params["ckpt_path"] = None