Commit 3d524fd3 authored by Aarni Koskela's avatar Aarni Koskela
Browse files

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

parent 8f6b24ce

modules/mac_specific.py

+4 −0
Original line number Diff line number Diff line
@@ -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:

CRA Git | Maintained and supported by SUSTech CRA and CCSE