auto enable TPU checkpoint save and checkpoint load using the proper wrappers

[patil-suraj]

I trained a model on TPU with google colab. When trying to load the checkpoint on GPU it gives following error

RuntimeError: Could not run 'aten::empty.memory_format' with arguments from the 'XLATensorId' backend. 'aten::empty.memory_format' is only available for these backends: [CUDATensorId, SparseCPUTensorId, VariableTensorId, CPUTensorId, MkldnnCPUTensorId, SparseCUDATensorId].

How to load the checkpoint saved on TPU with cpu/gpu ?

[williamFalcon]

ummm great point. @dlibenzi any docs on this?

[dlibenzi]

The recommended way to save checkpoint is to use the xm.save() API.
This is a thin wrapper around torch.save() that makes sure there is a symmetric graph execution on all cores:

https://github.com/pytorch/xla/blob/3518650335dba07820df8d5f6a33d1a969f78265/torch_xla/core/xla_model.py#L521

Saved tensors are PyTorch CPU tensor.
The load happen using vanilla torch.load() which sucks in PyTorch CPU tensor, then create your normal model, set the state dictionary, and then model.to(xla_device).

[williamFalcon]

i thought we had automated this.
this should be embedded in the auto checkpoint feature and load from checkpoint.

let’s turn this PR into those feature requests

[Borda]

will be fixed in #2726

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!