TPUBf16Precision

class lightning_fabric.plugins.precision.TPUBf16Precision

Bases: lightning_fabric.plugins.precision.tpu.TPUPrecision

Plugin that enables bfloats on TPUs.

convert_input(data)

Convert model inputs (forward) to the floating point precision type of this plugin.

This is a no-op for tensors that are not of floating-point type or already have the desired type.

Return type:

Tensor

teardown()

This method is called to teardown the training process.

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

Return type:

None