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
This commit is contained in:
Jeremy Reizenstein 2023-06-16 04:33:24 -07:00 committed by Facebook GitHub Bot
parent cd5db076d5
commit 009a3d3b3c

View File

@ -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: