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

pytorch3d/structures/meshes.py

@@ -329,11 +329,11 @@ class Meshes(object):
                 self._num_verts_per_mesh = torch.tensor(
                     [len(v) for v in self._verts_list], device=self.device
                 )
-                self._V = self._num_verts_per_mesh.max()
+                self._V = int(self._num_verts_per_mesh.max())
                 self._num_faces_per_mesh = torch.tensor(
                     [len(f) for f in self._faces_list], device=self.device
                 )
-                self._F = self._num_faces_per_mesh.max()
+                self._F = int(self._num_faces_per_mesh.max())
                 self.valid = torch.tensor(
                     [
                         len(v) > 0 and len(f) > 0
@@ -370,7 +370,7 @@ class Meshes(object):
                 # as long as the faces index correspond to the right vertices.
 
                 self.valid = self._num_faces_per_mesh > 0
-                self._F = self._num_faces_per_mesh.max()
+                self._F = int(self._num_faces_per_mesh.max())
                 if len(self._num_faces_per_mesh.unique()) == 1:
                     self.equisized = True
 

pytorch3d/structures/pointclouds.py

@@ -189,7 +189,7 @@ class Pointclouds(object):
                 num_points_per_cloud = torch.tensor(
                     [len(p) for p in self._points_list], device=self.device
                 )
-                self._P = num_points_per_cloud.max()
+                self._P = int(num_points_per_cloud.max())
                 self.valid = torch.tensor(
                     [len(p) > 0 for p in self._points_list],
                     dtype=torch.bool,

tests/test_point_mesh_distance.py

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