From d87e52ea2c112fd95f257dcd6a54a5db77a34fc3 Mon Sep 17 00:00:00 2001 From: Mel Massadian Date: Sun, 16 Feb 2025 02:52:26 +0100 Subject: [PATCH] =?UTF-8?q?fix:=20=F0=9F=90=9B=20stackImages=20move=20to?= =?UTF-8?q?=20device?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nodes/image_utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/nodes/image_utils.py b/nodes/image_utils.py index c0b683d..a2ccc3f 100644 --- a/nodes/image_utils.py +++ b/nodes/image_utils.py @@ -39,9 +39,13 @@ class MTB_StackImages: f"{'vertically' if vertical else 'horizontally'}" ) + target_device = tensors[0].device + normalized_tensors = [ - self.normalize_to_rgba(tensor) for tensor in tensors + self.normalize_to_rgba(tensor.to(target_device)) + for tensor in tensors ] + max_batch_size = max(tensor.shape[0] for tensor in normalized_tensors) normalized_tensors = [ self.duplicate_frames(tensor, max_batch_size)