From 8c3f093ecb6d38a4e4d0a9d331ee7bb3d7555a74 Mon Sep 17 00:00:00 2001 From: hturki Date: Thu, 1 Dec 2022 02:06:23 -0500 Subject: [PATCH] Fix device mismatch error when using masks (#1061) * Fix device mismatch error when using masks * Update pixel_samplers.py --- 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 b76d39f16..1066155c7 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], as_tuple=False) + nonzero_indices = torch.nonzero(batch["mask"][..., 0].to(device), as_tuple=False) chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch) indices = nonzero_indices[chosen_indices] else: