From f9a62a1d0f5b2aee7f98a10a486d14730e1979a5 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 5 Feb 2024 22:00:51 +0200
Subject: [PATCH] Update nodes.py

---
 nodes.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/nodes.py b/nodes.py
index 20af270..11803f3 100644
--- a/nodes.py
+++ b/nodes.py
@@ -3521,15 +3521,14 @@ class ImageUpscaleWithModelBatched:
         
         for start_idx in range(0, in_img.shape[0], per_batch):
             sub_images = upscale_model(in_img[start_idx:start_idx+per_batch])
-            t.append(sub_images)
+            t.append(sub_images.cpu())
             # Calculate the number of images processed in this batch
             batch_count = sub_images.shape[0]
             # Update the progress bar by the number of images processed in this batch
             pbar.update(batch_count)
         upscale_model.cpu()
         
-     
-        t = torch.cat(t, dim=0).permute(0, 2, 3, 1)
+        t = torch.cat(t, dim=0).permute(0, 2, 3, 1).cpu()
 
         return (t,)