Update train_util.py

gesen2egee · Feb 26, 2024 · 09d0696 · 09d0696
1 parent a649689
commit 09d0696
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/library/train_util.py b/library/train_util.py
@@ -1363,10 +1363,11 @@ def __getitem__(self, index):
                 image = self.image_transforms(img)  # -1.0~1.0のtorch.Tensorになる
 
             images.append(image)
-            latents_list.append(latents)
-            masks.append(torch.tensor(mask)) 
-
+            latents_list.append(latents) 
             target_size = (image.shape[2], image.shape[1]) if image is not None else (latents.shape[2] * 8, latents.shape[1] * 8)
+            if mask == None:
+                mask = np.full(target_shape, 255, np.uint8)            
+            masks.append(torch.tensor(mask))
 
             if not flipped:
                 crop_left_top = (crop_ltrb[0], crop_ltrb[1])