Fix device mismatch error when using masks (nerfstudio-project#1061)

* Fix device mismatch error when using masks * Update pixel_samplers.py
chris838 · Dec 1, 2022 · fd6b103 · fd6b103
1 parent d1bfc66
commit fd6b103
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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: