fix collision loss!!!!!!!!!!!!!!!!!!

losses/losses.py

@@ -4,9 +4,9 @@ import numpy as np
 # import tensorflow_graphics as tfg
 import scipy as sp
 import scipy.io as sio
+from scipy.spatial import KDTree
 
 from pytorch3d.structures import Meshes
-from scipy.spatial import cKDTree
 from torch_cluster import knn
 # from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
 # from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
@@ -33,15 +33,26 @@ def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
     batch_size = V.shape[0]
     vertices_count = V.shape[1]
     batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(V.shape[1])
-    batch_y = torch.arange(batch_size, device=B.get_device()).repeat_interleave(B.shape[1])
     V = V.reshape(-1, V.shape[-1])
     B = B.reshape(-1, B.shape[-1])
-    idx = knn(V, B.to(V.dtype), k=1, batch_x=batch_x, batch_y=batch_y)[1]
+    idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_x, batch_y=batch_x)[1]
     idx = (idx - batch_x * vertices_count).reshape(batch_size, vertices_count)
     return idx
 
+# def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
+#     # create an empty tensor to store the indices of the nearest neighbors
+#     batch_size = V.shape[0]
+#     vertices_count = V.shape[1]
+#     idx = torch.zeros((batch_size, vertices_count), dtype=torch.long, device=V.device)
+
+#     # loop over the batch dimension and find the nearest neighbor for each point in B
+#     for i in range(batch_size):
+#         tree = KDTree(B[i].cpu().numpy())
+#         _, nearest = tree.query(V[i].cpu().numpy(), k=1)
+#         idx[i] = torch.from_numpy(nearest)
+
+#     return idx
 
-# @tf.function
 def collision_loss(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr: float = .004):
     if B.shape[1] == 14475:
         F_norm_local = F_norm