Update train_util.py

gesen2egee · Feb 26, 2024 · 389f6cc · 389f6cc
1 parent 6164b4f
commit 389f6cc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/library/train_util.py b/library/train_util.py
@@ -1366,7 +1366,7 @@ def __getitem__(self, index):
             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 is None:
-                mask = np.full(target_shape, 1, np.uint8)            
+                mask = np.full(target_size, 1, np.uint8)            
             masks.append(torch.tensor(mask))
 
             if not flipped: