From 3e4586ede7047a5f50e35cde7928db27b3f25ac2 Mon Sep 17 00:00:00 2001 From: KimKyungPyo Date: Thu, 26 Mar 2026 17:46:46 +0900 Subject: [PATCH] Fix: properly handle batched tensors to prevent data loss --- comfy_extras/nodes_dataset.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index d0dab055e..d05c8c257 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -168,10 +168,15 @@ def save_images_to_folder(image_list, output_dir, prefix="image"): saved_files = [] if isinstance(image_list, torch.Tensor): - if image_list.dim() == 4: - image_list = [image_list[i] for i in range(image_list.shape[0])] + image_list = [image_list] + + normalized_images = [] + for img in image_list: + if isinstance(img, torch.Tensor) and img.dim() == 4: + normalized_images.extend([img[i] for i in range(img.shape[0])]) else: - image_list = [image_list] + normalized_images.append(img) + image_list = normalized_images for idx, img_tensor in enumerate(image_list): # Handle different tensor shapes