lightning_template.utils.callbacks.set_precision_and_cudnn_callback

Module Contents

Classes

SetPrecisionAndCudnnCallback

Set fp32 and fp16 precision and cudnn when training starts.

class lightning_template.utils.callbacks.set_precision_and_cudnn_callback.SetPrecisionAndCudnnCallback(float32_matmul_precision=None, allow_fp16_reduced_precision_reduction=None, deterministic_debug_mode=None, cudnn_enabled=None)

Bases: lightning.pytorch.Callback

Set fp32 and fp16 precision and cudnn when training starts.

For details, see https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html https://pytorch.org/docs/stable/generated/torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction.html https://pytorch.org/docs/stable/generated/torch.set_deterministic_debug_mode.html#torch.set_deterministic_debug_mode https://pytorch.org/docs/stable/backends.html#module-torch.backends.cudnn