def test_single_gpu():
    """Tests if device is set correctly when training and after teardown for single GPU strategy.

    Cannot run this test on MPS due to shared memory not allowing dedicated measurements of GPU memory utilization.

    """
    trainer = Trainer(accelerator="gpu", devices=1, fast_dev_run=True)
    # assert training strategy attributes for device setting
    assert isinstance(trainer.strategy, SingleDeviceStrategy)
    assert trainer.strategy.root_device == torch.device("cuda:0")

    model = BoringModelGPU()

    trainer.fit(model)

    # assert after training, model is moved to CPU and memory is deallocated
    assert model.device == torch.device("cpu")
    cuda_memory = torch.cuda.memory_allocated()
    assert cuda_memory < model.start_cuda_memory
Tests if device is set correctly when training and after teardown for single GPU strategy.

    Cannot run this test on MPS due to shared memory not allowing dedicated measurements of GPU memory utilization.

    
tests/tests_pytorch/strategies/test_single_device.py