add additional tf32 opt for cudnn (#2477) [skip ci]

2025-04-03 08:47:52 -04:00
parent 3877c5c69d
commit 5249e98058
1 changed files with 1 additions and 0 deletions
--- a/src/axolotl/utils/config/init.py
+++ b/src/axolotl/utils/config/init.py
@@ -78,6 +78,7 @@ def resolve_dtype(cfg):
        cfg.bf16 = False
    else:
        torch.backends.cuda.matmul.allow_tf32 = cfg.tf32 or False
+        torch.backends.cudnn.allow_tf32 = cfg.tf32 or False
        if cfg.bf16:
            cfg.fp16 = False