avoid torch.range

Summary: Avoid unintended use of torch.range.

Reviewed By: kjchalup

Differential Revision: D40341396

fbshipit-source-id: 108295983afdec0ca9e43178fef9c65695150bc1
This commit is contained in:
Jeremy Reizenstein 2022-10-13 14:37:46 -07:00 committed by Facebook GitHub Bot
parent 4d9215b3b4
commit f13086779d

View File

@ -207,7 +207,7 @@ class MultinomialRaysampler(torch.nn.Module):
n_rays_per_image,
) = _sample_cameras_and_masks(n_rays_total, cameras, mask)
else:
camera_ids = torch.range(0, len(cameras), dtype=torch.long)
camera_ids = torch.arange(len(cameras), dtype=torch.long)
batch_size = cameras.R.shape[0]
device = cameras.device
@ -438,7 +438,7 @@ class MonteCarloRaysampler(torch.nn.Module):
n_rays_per_image,
) = _sample_cameras_and_masks(self._n_rays_total, cameras, None)
else:
camera_ids = torch.range(0, len(cameras), dtype=torch.long)
camera_ids = torch.arange(len(cameras), dtype=torch.long)
n_rays_per_image = self._n_rays_per_image
batch_size = cameras.R.shape[0]