device_tpu#
The TPU device used for training.
Classes
- class composer.trainer.devices.device_tpu.DeviceTPU[source]#
Bases:
composer.trainer.devices.device.Device
An extension of
Device
for TPUs.When running on TPUVMs, you need to export PJRT_DEVICE=TPU. More details.