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

This commit is contained in:
Aarni Koskela 2023-07-12 15:17:13 +03:00
parent 8f6b24ce59
commit 3d524fd3f1

View File

@ -30,6 +30,10 @@ has_mps = check_for_mps()
def torch_mps_gc() -> None: def torch_mps_gc() -> None:
try: 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 from torch.mps import empty_cache
empty_cache() empty_cache()
except Exception: except Exception: