fix point_to_mesh distance issue (#592)

Signed-off-by: Clement Fuji Tsang <cfujitsang@nvidia.com>
NVIDIAGameWorks · Jul 16, 2022 · 207e187 · 207e187
1 parent 615fffc
commit 207e187
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/kaolin/metrics/trianglemesh.py b/kaolin/metrics/trianglemesh.py
@@ -113,9 +113,9 @@ class _UnbatchedTriangleDistanceCuda(torch.autograd.Function):
     def forward(ctx, points, face_vertices):
         num_points = points.shape[0]
         num_faces = face_vertices.shape[0]
-        min_dist = torch.zeros((num_points), device='cuda', dtype=points.dtype)
-        min_dist_idx = torch.zeros((num_points), device='cuda', dtype=torch.long)
-        dist_type = torch.zeros((num_points), device='cuda', dtype=torch.int32)
+        min_dist = torch.zeros((num_points), device=points.device, dtype=points.dtype)
+        min_dist_idx = torch.zeros((num_points), device=points.device, dtype=torch.long)
+        dist_type = torch.zeros((num_points), device=points.device, dtype=torch.int32)
         _C.metrics.unbatched_triangle_distance_forward_cuda(
             points, face_vertices, min_dist, min_dist_idx, dist_type)
         ctx.save_for_backward(points.contiguous(), face_vertices.contiguous(),