def test_resume_training_on_cpu(tmp_path):
    """Checks if training can be resumed from a saved checkpoint on CPU."""
    # Train a model on TPU
    model = BoringModel()
    trainer = Trainer(max_epochs=1, accelerator="tpu", devices="auto", default_root_dir=tmp_path)
    trainer.fit(model)

    if trainer.world_size != trainer.num_devices:
        # we're in multinode. unless the filesystem is shared, only the main node will have access to the checkpoint
        # since we cannot know this, the code below needs to be skipped
        return

    model_path = trainer.checkpoint_callback.best_model_path

    # Verify saved Tensors are on CPU
    ckpt = torch.load(model_path, weights_only=True)
    weight_tensor = list(ckpt["state_dict"].values())[0]
    assert weight_tensor.device == torch.device("cpu")

    # Verify that training is resumed on CPU
    trainer = Trainer(max_epochs=1, default_root_dir=tmp_path)
    trainer.fit(model, ckpt_path=model_path)
Checks if training can be resumed from a saved checkpoint on CPU.
tests/tests_pytorch/accelerators/test_xla.py