From 009a3d3b3c11c12b2099a6fd55b2ac50a8bb304f Mon Sep 17 00:00:00 2001 From: Jeremy Reizenstein Date: Fri, 16 Jun 2023 04:33:24 -0700 Subject: [PATCH] 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 | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/projects/nerf/nerf/raysampler.py b/projects/nerf/nerf/raysampler.py index a6b3c270..69e99b9a 100644 --- a/projects/nerf/nerf/raysampler.py +++ b/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: