Update train_util.py

gesen2egee · Feb 16, 2024 · 2c51ab1 · 2c51ab1
1 parent 5a506fd
commit 2c51ab1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/library/train_util.py b/library/train_util.py
@@ -1364,7 +1364,7 @@ def __getitem__(self, index):
 
             images.append(image)
             latents_list.append(latents)
-            if masks:
+            if 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)