XLAPrecision

class lightning.fabric.plugins.precision.XLAPrecision(precision)[source]

Bases: Precision

Plugin for training with XLA.

Parameters:

precision (Literal['32-true', '16-true', 'bf16-true']) – Full precision (32-true) or half precision (16-true, bf16-true).

Raises:

ValueError – If unsupported precision is provided.

optimizer_step(optimizer, **kwargs)[source]

Hook to run the optimizer step.

Return type:

Any

teardown()[source]

This method is called to teardown the training process.

It is the right place to release memory and free other resources.

Return type:

None