Fix deriving default map location when there is extra state

src/lightning/pytorch/CHANGELOG.md

@@ -122,6 +122,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed validation of parameters of `plugins.precision.MixedPrecisionPlugin` ([#17687](https://github.com/Lightning-AI/lightning/pull/17687))
 
 
+- Fixed deriving default map location in `LightningModule.load_from_checkpoint` when there is extra state ([#17812](https://github.com/Lightning-AI/lightning/pull/17812))
+
+
 ## [2.0.3] - 2023-06-07
 
 ### Changed

src/lightning/pytorch/core/saving.py

@@ -25,6 +25,7 @@
 from typing import Any, Callable, Dict, IO, Optional, Type, Union
 from warnings import warn
 
+import torch
 import yaml
 from lightning_utilities.core.apply_func import apply_to_collection
 
@@ -90,7 +91,8 @@ def _load_from_checkpoint(
  state_dict = checkpoint["state_dict"]
  if not state_dict:
  raise ValueError(f"The state dict in {checkpoint_path!r} contains no parameters.")
- map_location = list(state_dict.values())[0].device
+
+ map_location = next((t for t in state_dict.values() if isinstance(t, torch.Tensor)), torch.tensor(0)).device
  assert isinstance(storage, pl.LightningModule)
  return storage.to(map_location)
 

tests/tests_pytorch/core/test_saving.py

@@ -55,3 +55,19 @@ def test_load_from_checkpoint_map_location_cpu_to_gpu(tmp_path, map_location):
  create_boring_checkpoint(tmp_path, BoringModel(), accelerator="cpu")
  model = BoringModel.load_from_checkpoint(f"{tmp_path}/checkpoint.ckpt", map_location=map_location)
  assert model.device.type == "cuda"
+
+
+@RunIf(min_cuda_gpus=1)
+def test_load_from_checkpoint_default_map_location_extra_state(tmp_path):
+ """Test that the default map_location gets chosen based on the device of the saved tensors."""
+
+ class ExtraStateModel(BoringModel):
+ def get_extra_state(self):
+ return {"extra": "state"}
+
+ def set_extra_state(self, state):
+ pass
+
+ create_boring_checkpoint(tmp_path, ExtraStateModel(), accelerator="cuda")
+ model = ExtraStateModel.load_from_checkpoint(f"{tmp_path}/checkpoint.ckpt", map_location=None)
+ assert model.device.type == "cuda"