1 files changed, 17 insertions, 0 deletions
diff --git a/modules/devices.py b/modules/devices.py
index 30d30b99..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")
@@ -14,3 +16,18 @@ def get_optimal_device():
         return torch.device("mps")
 
     return cpu
+
+
+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")