Fix mask check in PatchPixelSampler (#1971)

* fix mask check in PatchPixelSampler

* change torch.Tensor to Tensor
This commit is contained in:
Lukas Hoellein authored and GitHub committed 2023-05-22 10:43:14 -07:00
1 parent 27e2c874cb
commit d8f2446122
1 file changed
+1 -1
+1 -1
View File
@@ -278,7 +278,7 @@ class PatchPixelSampler(PixelSampler):
mask: Optional[Tensor] = None,
device: Union[torch.device, str] = "cpu",
) -> Int[Tensor, "batch_size 3"]:
if mask:
if isinstance(mask, Tensor):
# Note: if there is a mask, sampling reduces back to uniform sampling
indices = super().sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device)
else: