From b70b51cc7248f57dcf16add86701139432c21e5b Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Mon, 12 Sep 2022 16:34:13 +0300
Subject:  Allow TF32 in CUDA for increased performance #279

---
 modules/devices.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

(limited to 'modules/devices.py')

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")
-- 
cgit v1.2.3