Allow TF32 in CUDA for increased performance #279

2 files changed, 21 insertions, 0 deletions

diff --git a/modules/devices.py b/modules/devices.py
index f88e807e..a93a245b 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -1,6 +1,8 @@
 import torch
 
 # has_mps is only available in nightly pytorch (for now), `getattr` for compatibility
+from modules import errors
+
 has_mps = getattr(torch, 'has_mps', False)
 
 cpu = torch.device("cpu")
@@ -20,3 +22,12 @@ def torch_gc():
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
+
+
+def enable_tf32():
+    if torch.cuda.is_available():
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+
+
+errors.run(enable_tf32, "Enabling TF32")
diff --git a/modules/errors.py b/modules/errors.py
new file mode 100644
index 00000000..372dc51a
--- /dev/null
+++ b/modules/errors.py
@@ -0,0 +1,10 @@
+import sys

+import traceback

+

+

+def run(code, task):

+    try:

+        code()

+    except Exception as e:

+        print(f"{task}: {type(e).__name__}", file=sys.stderr)

+        print(traceback.format_exc(), file=sys.stderr)