load_from_checkpoint does not work as expected

[liujuncn]

Bug description

I want to load a trained checkpoint to "gpu" in colab, but it seems that load_from_checkpoint loads two copies, and the device of the model is "cpu".

The memory of both host and gpu is occupied.

If i use:
model.to(torch.device('cuda'))

subsequent model predict will report the following error:
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)

How to solve this problem？

What version are you seeing the problem on?

2.0+

How to reproduce the bug

No response

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

[ryan597]

Currently the object is created and returned without using map_location at all.
The checkpoint is mapped, but this doesn't seem to affect the model created and returned later.

https://github.com/Lightning-AI/lightning/blob/fd4697c62c059fc7b9946e84d91625ecb6efdbe5/src/lightning/pytorch/core/saving.py#L51-L92

I think this works to fix it

import torch
...
def _load_from_checkpoint(
...
    if issubclass(cls, pl.LightningModule):
        storage = _load_state(cls, checkpoint, strict=strict, **kwargs)
        restore_location = torch.serialization._get_restore_location(map_location)

        if isinstance(map_location, dict):
            return restore_location(storage, map_location.get(str(storage.device)))

        return restore_location(storage, map_location)

[ryan597]

@ryan597 Yes, it can be loaded to cuda device, but still report the error when predicting:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)

If using cpu to predict, it works but slowly.

Have you moved the data to the gpu?

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)

I get that error when the inputs are not moved to the gpu, see below

model = DemoModel.load_from_checkpoint("./checkpoints/last.ckpt", map_location=torch.device('cuda'))

model.setup()
data = model.test_dataloader()

with torch.no_grad():
    for i, [inp, out] in enumerate(data):
        inp = inp.to('cuda')  # also move inputs to gpu
        model(inp)

Note there is a slightly different error if the model is on the cpu and the data on the gpu (swaps cuda:0 and cpu)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)

[liujuncn]

yes, you are right. I forget to move the data to gpu.