`load_from_checkpoint` returns `None`

[shmuelyo]

Bug description

On calling load_from_checkpoint i get 'None'.
I tried debugging the issue and saw that the lightning module class is initialized successfully in _load_from_checkpoint (line 91 in saving.py), but at the end return model.to(device) returns None (line 99 in saving.py).
AFAIK this is the expected behavior of pytorch's to, which is why i believe this is a bug.
And just FYI, loading the model works fine when i initialize the lightning module class and call pytorch's load_state_dict.
I'm using pytorch-lightning 2.5.0.post0 with python 3.12 and pytorch 2.5.1.

What version are you seeing the problem on?

master

How to reproduce the bug

Error messages and logs

None

Environment

Current environment

* CUDA:
        - GPU:
                - Tesla T4
        - available:         True
        - version:           12.4
* Lightning:
        - lightning-utilities: 0.11.9
        - pytorch-lightning: 2.5.0.post0
        - torch:             2.5.1
        - torchmetrics:      1.6.1
        - torchvision:       0.20.1
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                - ELF
        - processor:         x86_64
        - python:            3.12.3

More info

No response

[GdoongMathew]

Hi @shmuelyo, I think torch always return the original module instance instead of None, when it comes to the .to(device, dtype) method. https://github.com/pytorch/pytorch/blob/a8d6afb511a69687bbb2b7e88a3cf67917e1697e/torch/nn/modules/module.py#L1031

[matsumotosan]

@shmuelyo I just tested with the script below and load_from_checkpoint returns a BoringModel as expected. Is this different from what you are seeing?

import os
import tempfile

from lightning import Trainer
from lightning.pytorch.demos import BoringModel

model = BoringModel()

with tempfile.TemporaryDirectory() as tempdir:
    trainer = Trainer(default_root_dir=tempdir, max_epochs=5)
    trainer.fit(model)
    ckpt_path = os.listdir(f"{tempdir}/lightning_logs/version_0/checkpoints")[0]
    model_from_ckpt = BoringModel.load_from_checkpoint(f"{tempdir}/lightning_logs/version_0/checkpoints/{ckpt_path}")
    print(type(model))

[Borda]

Great work @matsumotosan! Thank you 👍