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

pre_dispatch()

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

Return type

None

save_checkpoint(checkpoint, filepath)

Save model/training states as a checkpoint file through state-dump and file-write.

Parameters

• checkpoint (Dict[str, Any]) – dict containing model and trainer state

• filepath (str) – write-target file’s path

Return type

None

teardown()

This method is called to teardown the training process. It is the right place to release memory and free other resources.

Return type

None