Fix issue with not freeing enough memory when sampling.

be9c5e25 · comfyanonymous · ac0758a1 · be9c5e25
Commit be9c5e25 authored Aug 17, 2023 by comfyanonymous
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

comfy/sample.py comfy/sample.py +1 -1

No files found.
--- a/comfy/sample.py
+++ b/comfy/sample.py
@@ -78,7 +78,7 @@ def sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative

    real_model = None
    models = get_additional_models(positive, negative)
-    comfy.model_management.load_models_gpu([model] + models, comfy.model_management.batch_area_memory(noise.shape[2] * noise.shape[3]))
+    comfy.model_management.load_models_gpu([model] + models, comfy.model_management.batch_area_memory(noise.shape[0] * noise.shape[2] * noise.shape[3]))
    real_model = model.model

    noise = noise.to(device)