From aaff1b478cc81233126375a2222033b42af1a9c3 Mon Sep 17 00:00:00 2001 From: Tao Date: Wed, 14 Dec 2022 02:16:34 +0800 Subject: [PATCH] Fix slow training issue when using mask (#1114) --- 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 1066155c76..b76d39f16a 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -38,7 +38,7 @@ def collate_image_dataset_batch(batch: Dict, num_rays_per_batch: int, keep_full_ # only sample within the mask, if the mask is in the batch if "mask" in batch: - nonzero_indices = torch.nonzero(batch["mask"][..., 0].to(device), as_tuple=False) + nonzero_indices = torch.nonzero(batch["mask"][..., 0], as_tuple=False) chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch) indices = nonzero_indices[chosen_indices] else: