From 3d524fd3f1bdb17946bf6fa8a3cdf7b10859c495 Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Wed, 12 Jul 2023 15:17:13 +0300
Subject: Don't do MPS GC when there's a latent that could still be sampled

---
 modules/mac_specific.py | 4 ++++
 1 file changed, 4 insertions(+)

(limited to 'modules/mac_specific.py')

diff --git a/modules/mac_specific.py b/modules/mac_specific.py
index 328b5973..9ceb43ba 100644
--- a/modules/mac_specific.py
+++ b/modules/mac_specific.py
@@ -30,6 +30,10 @@ has_mps = check_for_mps()
 
 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:
-- 
cgit v1.2.3