XLAAccelerator¶

class lightning.pytorch.accelerators.XLAAccelerator(*args, **kwargs)[source]¶

Bases: Accelerator, XLAAccelerator

Accelerator for XLA devices, normally TPUs.

Warning

Use of this accelerator beyond import and instantiation is experimental.

get_device_stats(device)[source]¶

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.