DI-engine/ding/torch_utils/backend_helper.py

import torch


def enable_tf32() -> None:
    """
    Overview:
        Enable tf32 on matmul and cudnn for faster computation. This only works on Ampere GPU devices. \
        For detailed information, please refer to: \
        https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices.
    """
    torch.backends.cuda.matmul.allow_tf32 = True  # allow tf32 on matmul
    torch.backends.cudnn.allow_tf32 = True  # allow tf32 on cudnn