From 805e6e702d328bf6fc2bf9721cbb0a4eedf03a35 Mon Sep 17 00:00:00 2001 From: anc2001 <71151378+anc2001@users.noreply.github.com> Date: Thu, 14 Sep 2023 09:11:15 -0700 Subject: [PATCH] fix pixel sampling with masks issue when data is list (#2369) Co-authored-by: Brent Yi --- nerfstudio/data/pixel_samplers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index b852c31b1..463d08294 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -215,7 +215,7 @@ class PixelSampler: num_rays_in_batch = num_rays_per_batch - (num_images - 1) * num_rays_in_batch indices = self.sample_method( - num_rays_in_batch, 1, image_height, image_width, mask=batch["mask"][i], device=device + num_rays_in_batch, 1, image_height, image_width, mask=batch["mask"][i].unsqueeze(0), device=device ) indices[:, 0] = i all_indices.append(indices)