projects/nerf subsampling fix for newish pytorch #1441

Summary: Fix for https://github.com/facebookresearch/pytorch3d/issues/1441 where we were indexing with a tensor on the wrong device.

Reviewed By: shapovalov

Differential Revision: D46276449

fbshipit-source-id: 7750ed45ffecefa5d291fd1eadfe515310c2cf0d

projects/nerf/nerf/raysampler.py

@@ -330,9 +330,9 @@ class NeRFRaysampler(torch.nn.Module):
 
             if self.training:
                 # During training we randomly subsample rays.
-                sel_rays = torch.randperm(n_pixels, device=device)[
-                    : self._mc_raysampler._n_rays_per_image
-                ]
+                sel_rays = torch.randperm(
+                    n_pixels, device=full_ray_bundle.lengths.device
+                )[: self._mc_raysampler._n_rays_per_image]
             else:
                 # In case we test, we take only the requested chunk.
                 if chunksize is None: