`Trainer`'s `.init_module()` context does not initialize model on target device

[jin-zhe]

Bug description

I refer to the documentation on https://lightning.ai/docs/pytorch/stable/advanced/model_init.html which states "you can force PyTorch to create the model directly on the target device" when using the .init_module() context. However I have verified across different GPU machines that this is not the case. A simple code is provided below which prints out the model's device after initialization under the context. It always prints 'cpu'.

What version are you seeing the problem on?

v2.4

How to reproduce the bug

from torch import nn, optim
from pytorch_lightning import Trainer, LightningModule


class LitAutoEncoder(LightningModule):
  '''
  Model taken from https://lightning.ai/docs/pytorch/stable/starter/introduction.html
  Details unimportant
  '''
  def __init__(self):
    super().__init__()
    self.encoder = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
    self.decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))

  def training_step(self, batch, batch_idx):
    x, _ = batch
    x = x.view(x.size(0), -1)
    z = self.encoder(x)
    x_hat = self.decoder(z)
    loss = nn.functional.mse_loss(x_hat, x)
    return loss

  def configure_optimizers(self):
    optimizer = optim.Adam(self.parameters(), lr=1e-3)
    return optimizer

trainer = Trainer(accelerator='gpu', devices=[0])
with trainer.init_module():
  model = LitAutoEncoder()
  print(model.device) # => cpu

Error messages and logs

# Error messages and logs here please

Environment

Current environment

` * CUDA: - GPU: - Tesla V100-SXM2-32GB-LS - Tesla V100-SXM2-32GB - Tesla V100-SXM2-32GB-LS - Tesla V100-SXM2-32GB-LS - Tesla V100-SXM2-32GB-LS - Tesla V100-SXM2-32GB-LS - Tesla V100-SXM2-32GB-LS - Tesla V100-SXM2-32GB-LS - available: True - version: 11.8 * Lightning: - lightning: 2.4.0 - lightning-utilities: 0.11.6 - open-clip-torch: 2.26.1 - pytorch-lightning: 2.4.0 - torch: 2.1.0 - torchaudio: 2.1.0 - torchmetrics: 1.4.0.post0 - torchvision: 0.16.0 * Packages: - aiohappyeyeballs: 2.3.5 - aiohttp: 3.10.3 - aiosignal: 1.3.1 - altair: 5.4.1 - antlr4-python3-runtime: 4.9.3 - appdirs: 1.4.4 - asttokens: 2.4.1 - async-timeout: 4.0.3 - attrs: 24.2.0 - autocommand: 2.2.2 - backports.tarfile: 1.2.0 - blinker: 1.8.2 - brotli: 1.1.0 - cachetools: 5.5.0 - certifi: 2024.8.30 - cffi: 1.17.0 - charset-normalizer: 3.3.2 - click: 8.1.7 - colorama: 0.4.6 - comm: 0.2.2 - datasets: 2.20.0 - debugpy: 1.8.5 - decorator: 5.1.1 - dill: 0.3.8 - docker-pycreds: 0.4.0 - einops: 0.8.0 - exceptiongroup: 1.2.2 - executing: 2.1.0 - filelock: 3.15.4 - frozenlist: 1.4.1 - fsspec: 2024.5.0 - ftfy: 6.2.3 - gitdb: 4.0.11 - gitpython: 3.1.43 - gmpy2: 2.1.5 - h2: 4.1.0 - hpack: 4.0.0 - huggingface-hub: 0.24.5 - hyperframe: 6.0.1 - idna: 3.7 - importlib-metadata: 7.2.1 - importlib-resources: 6.4.5 - inflect: 7.3.1 - ipykernel: 6.29.5 - ipython: 8.27.0 - ipywidgets: 8.1.5 - jaraco.context: 5.3.0 - jaraco.functools: 4.0.1 - jaraco.text: 3.12.1 - jedi: 0.19.1 - jinja2: 3.1.4 - jsonlines: 4.0.0 - jsonschema: 4.23.0 - jsonschema-specifications: 2023.12.1 - jupyter-client: 8.6.3 - jupyter-core: 5.7.2 - jupyterlab-widgets: 3.0.13 - lightning: 2.4.0 - lightning-utilities: 0.11.6 - markdown-it-py: 3.0.0 - markupsafe: 2.1.5 - matplotlib-inline: 0.1.7 - mdurl: 0.1.2 - more-itertools: 10.3.0 - mpmath: 1.3.0 - multidict: 6.0.5 - multiprocess: 0.70.16 - narwhals: 1.8.2 - nest-asyncio: 1.6.0 - networkx: 3.3 - numpy: 1.26.4 - omegaconf: 2.3.0 - open-clip-torch: 2.26.1 - opencv-python: 4.10.0 - opencv-python-headless: 4.10.0 - ordered-set: 4.1.0 - packaging: 24.1 - pandas: 2.2.2 - parso: 0.8.4 - pathtools: 0.1.2 - pexpect: 4.9.0 - pickleshare: 0.7.5 - pillow: 10.4.0 - pip: 24.2 - pkgutil-resolve-name: 1.3.10 - platformdirs: 4.3.6 - prompt-toolkit: 3.0.47 - protobuf: 4.25.3 - psutil: 6.0.0 - ptyprocess: 0.7.0 - pure-eval: 0.2.3 - pyarrow: 17.0.0 - pyarrow-hotfix: 0.6 - pycparser: 2.22 - pydeck: 0.8.0b4 - pygments: 2.18.0 - pysocks: 1.7.1 - python-dateutil: 2.9.0 - pytorch-lightning: 2.4.0 - pytz: 2024.1 - pyyaml: 6.0.2 - pyzmq: 26.2.0 - referencing: 0.35.1 - regex: 2024.7.24 - requests: 2.32.3 - rich: 13.8.1 - rpds-py: 0.20.0 - safetensors: 0.4.4 - sentry-sdk: 2.12.0 - setproctitle: 1.3.3 - setuptools: 72.1.0 - six: 1.16.0 - smmap: 5.0.0 - stack-data: 0.6.2 - streamlit: 1.38.0 - sympy: 1.13.2 - tenacity: 8.5.0 - timm: 1.0.8 - tokenizers: 0.19.1 - toml: 0.10.2 - tomli: 2.0.1 - torch: 2.1.0 - torchaudio: 2.1.0 - torchmetrics: 1.4.0.post0 - torchvision: 0.16.0 - tornado: 6.4.1 - tqdm: 4.66.5 - traitlets: 5.14.3 - transformers: 4.44.2 - triton: 2.1.0 - typeguard: 4.3.0 - typing-extensions: 4.12.2 - tzdata: 2024.1 - tzlocal: 5.2 - urllib3: 2.2.2 - validators: 0.34.0 - wandb: 0.16.6 - watchdog: 4.0.1 - wcwidth: 0.2.13 - wheel: 0.44.0 - widgetsnbextension: 4.0.13 - xformers: 0.0.22.post7 - xxhash: 3.4.1 - yarl: 1.9.4 - zipp: 3.20.2 - zstandard: 0.23.0 * System: - OS: Linux - architecture: - 64bit - ELF - processor: x86_64 - python: 3.10.14 - release: 4.15.0-55-generic - version: #60-Ubuntu SMP Tue Jul 2 18:22:20 UTC 2019 ```

#- How you installed Lightning(conda, pip, source): conda

More info

No response

[nocoding03]

By default, PyTorch creates the model on the CPU, unless you explicitly move it to a different device (e.g. via .to(device)). As a result, even with trainer.init_module(), the model will still be initialized on the CPU first, and you will need to manually move the model to the GPU device.

Solution:
You can explicitly move the model to the target device in the context of with trainer.init_module():
with trainer.init_module():
model = LitAutoEncoder()
model = model.to('cuda') # Explicitly move the model to the GPU
print(next(model.parameters()).device) # The device on which the model is printed