1 files changed, 10 insertions, 0 deletions
diff --git a/modules/devices.py b/modules/devices.py
index b7899632..0158b11f 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -1,3 +1,5 @@
+import contextlib
+
 import torch
 
 from modules import errors
@@ -56,3 +58,11 @@ def randn_without_seed(shape):
 
     return torch.randn(shape, device=device)
 
+
+def autocast():
+    from modules import shared
+
+    if dtype == torch.float32 or shared.cmd_opts.precision == "full":
+        return contextlib.nullcontext()
+
+    return torch.autocast("cuda")