XLAAccelerator

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

Bases: Accelerator, XLAAccelerator

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.