From d875f5f5b5aae26fc839b42df765458d1e55f3eb Mon Sep 17 00:00:00 2001 From: Yuantao Chen <44546388+Tao-11-chen@users.noreply.github.com> Date: Wed, 26 Apr 2023 09:32:20 +0800 Subject: [PATCH] fix_mask (#1819) fix_nerfacto_mask Co-authored-by: machenmusik --- nerfstudio/data/pixel_samplers.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 244498df8..b58d99fe3 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -96,6 +96,7 @@ class PixelSampler: # pylint: disable=too-few-public-methods indices = self.sample_method(num_rays_per_batch, num_images, image_height, image_width, device=device) c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) + c, y, x = c.cpu(), y.cpu(), x.cpu() collated_batch = { key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None }