fix bug in getting mask_pixel_values for non-normal training settings (#252)

This commit is contained in:
Tianhao-Qi
2025-07-18 14:13:44 +08:00
committed by GitHub
parent 9436ff2f6e
commit 4f14db1269
+1 -1
View File
@@ -498,7 +498,7 @@ class ImageVideoDataset(Dataset):
if self.enable_inpaint and not self.enable_bucket:
mask = get_random_mask(pixel_values.size())
mask_pixel_values = pixel_values * (1 - mask) + torch.ones_like(pixel_values) * -1 * mask
mask_pixel_values = pixel_values * (1 - mask) + torch.ones_like(pixel_values) * mask
sample["mask_pixel_values"] = mask_pixel_values
sample["mask"] = mask