Removed for loop in sliced wasserstein distance

aidos-lab · Jan 30, 2024 · aba1474 · aba1474
1 parent 5426f42
commit aba1474
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_topological/nn/sliced_wasserstein_distance.py b/torch_topological/nn/sliced_wasserstein_distance.py
@@ -83,8 +83,8 @@ def forward(self, X, Y):
             diag = torch.tensor([0.5, 0.5], dtype=torch.float32)
 
             # Project both the diagrams onto the diagonals.
-            D1_diag = torch.vstack([torch.sum(x) * diag for x in D1])
-            D2_diag = torch.vstack([torch.sum(x) * diag for x in D2])
+            D1_diag = torch.sum(D1, dim=1, keepdim=True) * diag
+            D2_diag = torch.sum(D2, dim=1, keepdim=True) * diag
 
             cost = 0.0