feat: allow tf32 on cudnn

jolibrain · Dec 10, 2024 · 367cd91 · 367cd91
1 parent f727847
commit 367cd91
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/train.py b/train.py
@@ -441,6 +441,7 @@ def launch_training(opt):
 
     if opt.with_tf32:
         torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
 
     opt.use_cuda = torch.cuda.is_available() and opt.gpu_ids and opt.gpu_ids[0] >= 0
     if opt.use_cuda: