SingleTPUPlugin

class pytorch_lightning.plugins.training_type.SingleTPUPlugin(device, debug=False)

Bases: pytorch_lightning.plugins.training_type.single_device.SingleDevicePlugin

Plugin for training on a single TPU device.

model_to_device()

Moves the model to the correct device

Return type

None

on_save(checkpoint)

Move XLA tensors to CPU before saving Recommended on XLA Guide: https://github.com/pytorch/xla/blob/master/API_GUIDE.md#saving-and-loading-xla-tensors

Return type

dict

pre_dispatch()

Hook to do something before the training/evaluation/prediction starts.

Return type

None