def test_setup_hook_device_and_layers(tmp_path, accelerator):
    """Test `LightningModule.device` access and creation of layers in `LightningModule.setup` hook."""
    expected_device = torch.device(accelerator, 0)

    class TestModel(BoringModel):
        def setup(self, stage: str) -> None:
            # The `self.device` attribute already points to what device the model will land on
            assert self.device == expected_device
            # However, the model parameters have not yet been moved to that device
            assert self.layer.weight.device == torch.device("cpu")
            # Can create new layers in this hook (on CPU)
            self.new_layer = torch.nn.Linear(2, 2)
            assert self.new_layer.weight.device == torch.device("cpu")

        def training_step(self, batch, batch_idx):
            output = self.layer(batch)
            # will crash if not moved to correct device
            output = self.new_layer(output)
            loss = self.loss(output)
            return {"loss": loss}

    # fake data
    train_data = torch.utils.data.DataLoader(RandomDataset(32, 64))

    # model
    model = TestModel()
    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True, accelerator=accelerator, devices=1)
    trainer.fit(model, train_data)
Test `LightningModule.device` access and creation of layers in `LightningModule.setup` hook.
tests/tests_pytorch/trainer/test_trainer.py