fixed a bug in the normalized laplacian implementation and a bug in b…

…atch for sparse graph creation
shaham-lab · Nov 25, 2023 · da19231 · da19231
1 parent 53a5c8b
commit da19231
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/spectralnet/_losses/_spectralnet_loss.py b/src/spectralnet/_losses/_spectralnet_loss.py
@@ -25,7 +25,7 @@ def forward(
         m = Y.size(0)
         if is_normalized:
             D = torch.sum(W, dim=1)
-            Y = Y / D[:, None]
+            Y = Y / torch.sqrt(D)[:, None]
 
         Dy = torch.cdist(Y, Y)
         loss = torch.sum(W * Dy.pow(2)) / (2 * m)

diff --git a/src/spectralnet/_utils.py b/src/spectralnet/_utils.py
@@ -62,7 +62,7 @@ def make_batch_for_sparse_grapsh(batch_x: torch.Tensor) -> torch.Tensor:
         x = x.detach().cpu().numpy()
         nn_indices = u.get_nns_by_vector(x, n_neighbors)
         nn_tensors = [u.get_item_vector(i) for i in nn_indices[1:]]
-        nn_tensors = torch.tensor(nn_tensors)
+        nn_tensors = torch.tensor(nn_tensors, device=batch_x.device)
         new_batch_x = torch.cat((new_batch_x, nn_tensors))
 
     return new_batch_x