Set mass of set_a to ones

fwilliams · Feb 20, 2019 · a1ec9e0 · a1ec9e0
1 parent 399319c
commit a1ec9e0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/sinkhorn_loss_unbalanced/main.py b/examples/sinkhorn_loss_unbalanced/main.py
@@ -36,11 +36,11 @@
     print(set_b)
 
     # Condition P*1 = a and P^T*1 = b
-    a = torch.rand(set_a.shape[0:2],
+    a = torch.ones(set_a.shape[0:2],
             requires_grad=False,
             device=set_a.device)
 
-    b = torch.rand(set_b.shape[0:2],
+    b = torch.ones(set_b.shape[0:2],
             requires_grad=False,
             device=set_b.device)
     # Have the same total mass than set_a