def test_get_model_gpu(tmp_path, accelerator):
    """Tests that `trainer.lightning_module` extracts the model correctly when using GPU."""
    model = TrainerGetModel()

    limit_train_batches = 2
    trainer = Trainer(
        default_root_dir=tmp_path,
        limit_train_batches=limit_train_batches,
        limit_val_batches=2,
        max_epochs=1,
        accelerator=accelerator,
        devices=1,
    )
    trainer.fit(model)
Tests that `trainer.lightning_module` extracts the model correctly when using GPU.
tests/tests_pytorch/trainer/properties/test_get_model.py