From 7f1e63aed1252ba8145d4a66ce2272331d60cdae Mon Sep 17 00:00:00 2001 From: Jeremy Reizenstein Date: Wed, 10 Jun 2020 14:11:10 -0700 Subject: [PATCH] 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 | 6 +++--- pytorch3d/structures/pointclouds.py | 2 +- tests/test_point_mesh_distance.py | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/pytorch3d/structures/meshes.py b/pytorch3d/structures/meshes.py index 32c2b692..0f0b963a 100644 --- a/pytorch3d/structures/meshes.py +++ b/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 diff --git a/pytorch3d/structures/pointclouds.py b/pytorch3d/structures/pointclouds.py index 687de83c..146fb60b 100644 --- a/pytorch3d/structures/pointclouds.py +++ b/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, diff --git a/tests/test_point_mesh_distance.py b/tests/test_point_mesh_distance.py index 9fe18293..d914dcb8 100644 --- a/tests/test_point_mesh_distance.py +++ b/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]