From 10b5ca25418a670d0c4f27fe5846cd1eb0e12d43 Mon Sep 17 00:00:00 2001
From: lllyasviel <lyuminzhang@outlook.com>
Date: Fri, 8 Mar 2024 00:40:33 -0800
Subject: [PATCH] vae already sliced in inner loop

---
 modules/processing.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/modules/processing.py b/modules/processing.py
index 4528866c..64e564e0 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -633,10 +633,10 @@ class DecodedSamples(list):
 
 def decode_latent_batch(model, batch, target_device=None, check_for_nans=False):
     samples = DecodedSamples()
+    samples_pytorch = decode_first_stage(model, batch).to(target_device)
 
-    for i in range(batch.shape[0]):
-        sample = decode_first_stage(model, batch[i:i + 1])[0]
-        samples.append(sample.to(target_device))
+    for x in samples_pytorch:
+        samples.append(x)
 
     return samples