Fix device mismatch error when using masks

[hturki]

When using masks, I get the following error:

  File "/home/hturki/miniconda3/envs/my-end/lib/python3.9/site-packages/nerfstudio/data/pixel_samplers.py", line 51, in <dictcomp>
    collated_batch = {key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None}
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

printing the keys and devices in my batch, I see:

image_idx cuda:0
image cpu
mask cuda:0

Not sure if this fix is the most elegant solution

[tancik]

Does it by any chance work if you instead do
nonzero_indices = torch.nonzero(batch["mask"][..., 0].to(device), as_tuple=False)
This puts it on the gpu earlier which might be a little faster

[hturki]

that seems to work too, but FYI 'device' in this case seems to be cpu (sorry, had a typo in the original PR description)

[tancik]

LGTM