Update train_util.py

gesen2egee · Feb 16, 2024 · c41ce75 · c41ce75
1 parent 2c51ab1
commit c41ce75
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/library/train_util.py b/library/train_util.py
@@ -1364,8 +1364,7 @@ def __getitem__(self, index):
 
             images.append(image)
             latents_list.append(latents)
-            if mask:
-                masks.append(torch.tensor(mask)) 
+            masks.append(torch.tensor(mask)) 
 
             target_size = (image.shape[2], image.shape[1]) if image is not None else (latents.shape[2] * 8, latents.shape[1] * 8)