use torch.diagonal to set the diagonal of the laplacian (#643)

Signed-off-by: Clement Fuji Tsang <cfujitsang@nvidia.com> Signed-off-by: Clement Fuji Tsang <cfujitsang@nvidia.com>
NVIDIAGameWorks · Nov 2, 2022 · b47f216 · b47f216
1 parent bcce7a9
commit b47f216
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/kaolin/ops/mesh/mesh.py b/kaolin/ops/mesh/mesh.py
@@ -113,8 +113,7 @@ def uniform_laplacian(num_vertices, faces):
 
     L = torch.div(dense_adjacency, num_neighbour)
 
-    mask = torch.eye(num_vertices, num_vertices, device=faces.device, dtype=torch.bool)
-    L = L.masked_fill_(mask, -1)
+    torch.diagonal(L)[:] = -1
 
     # Fill NaN value with 0
     L[torch.isnan(L)] = 0