Don't do MPS GC when there's a latent that could still be sampled

Asdasdasda223 · Jul 12, 2023 · 3d524fd · 3d524fd
1 parent 8f6b24c
commit 3d524fd
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/modules/mac_specific.py b/modules/mac_specific.py
@@ -30,6 +30,10 @@ def check_for_mps() -> bool:
 
 def torch_mps_gc() -> None:
     try:
+        from modules.shared import state
+        if state.current_latent is not None:
+            log.debug("`current_latent` is set, skipping MPS garbage collection")
+            return
         from torch.mps import empty_cache
         empty_cache()
     except Exception: