XLAAccelerator

class lightning.pytorch.accelerators.XLAAccelerator(*args, **kwargs)

Accelerator for XLA devices, normally TPUs.

Warning

Use of this accelerator beyond import and instantiation is experimental.

get_device_stats(device)

Gets stats for the given XLA device.

Parameters:: device (Union[device, str, int]) – XLA device for which to get stats
Return type:: dict[str, Any]
Returns:: A dictionary mapping the metrics (free memory and peak memory) to their values.

Was this page helpful?