fix: 🐛 tensor2pil

melMass · Jul 21, 2023 · 8a59508 · 8a59508
1 parent aa551eb
commit 8a59508
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/utils.py b/utils.py
@@ -53,7 +53,8 @@ def tensor2pil(image: torch.Tensor) -> List[Image.Image]:
 
     if batch_count > 1:
         out = []
-        out.extend([tensor2pil(image[i]) for i in range(batch_count)])
+        for i in range(batch_count):
+            out.extend(tensor2pil(image[i]))
         return out
 
     return [