force tf32 to be enabled (torch 1.12 disables it)

2022-07-16 13:59:07 -06:00 · 2022-07-16 13:59:07 -06:00 · a073fbfcb8
commit a073fbfcb8
parent 438dcaccc5
1 changed files with 1 additions and 0 deletions
--- a/codes/train.py
+++ b/codes/train.py
@ -94,6 +94,7 @@ class Trainer:
        util.set_random_seed(seed)

        torch.backends.cudnn.benchmark = opt_get(opt, ['cuda_benchmarking_enabled'], True)
+        torch.backends.cuda.matmul.allow_tf32 = True
        # torch.backends.cudnn.deterministic = True
        if opt_get(opt, ['anomaly_detection'], False):
            torch.autograd.set_detect_anomaly(True)