Take care with single integers on gpu

Summary:
Pytorch seems to be becoming stricter about integer tensors of shape `(1,)` on GPU and not allowing them to be used as `int`s. For example the following no longer works on pytorch master,
    foo = torch.tensor([3, 5, 3], device="cuda:0")
    torch.arange(10) + foo[0]
because this is the sum of tensors on different devices.

Here fix tests which recently broke because of this.

Reviewed By: nikhilaravi

Differential Revision: D21929745

fbshipit-source-id: 25374f70468d1c895372766f1a9dd61df0833957
This commit is contained in:
Jeremy Reizenstein
2020-06-10 14:11:10 -07:00
committed by Facebook GitHub Bot
parent d0e7426a06
commit 7f1e63aed1
3 changed files with 7 additions and 7 deletions

View File

@@ -385,7 +385,7 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
start = edges_first_idx[i]
end = edges_first_idx[i + 1] if i < N - 1 else edges_packed.shape[0]
min_idx = idx_cuda.cpu()[start:end] - points_first_idx[i]
min_idx = idx_cuda.cpu()[start:end] - points_first_idx[i].cpu()
iidx = torch.arange(edges.shape[0], device=device)
min_dist = dists_temp[iidx, min_idx]
@@ -583,7 +583,7 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
start = points_first_idx[i]
end = points_first_idx[i + 1] if i < N - 1 else points_packed.shape[0]
min_idx = idx_cuda.cpu()[start:end] - faces_first_idx[i]
min_idx = idx_cuda.cpu()[start:end] - faces_first_idx[i].cpu()
iidx = torch.arange(points.shape[0], device=device)
min_dist = dists_temp[iidx, min_idx]
@@ -666,7 +666,7 @@ class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
start = faces_first_idx[i]
end = faces_first_idx[i + 1] if i < N - 1 else faces_packed.shape[0]
min_idx = idx_cuda.cpu()[start:end] - points_first_idx[i]
min_idx = idx_cuda.cpu()[start:end] - points_first_idx[i].cpu()
iidx = torch.arange(tris.shape[0], device=device)
min_dist = dists_temp[iidx, min_idx]