pep 8

atong01 · Nov 23, 2023 · d700429 · d700429
1 parent 2672db1
commit d700429
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/tests/test_models.py b/tests/test_models.py
@@ -1,4 +1,5 @@
 import pytest
+
 from torchcfm.models import MLP
 from torchcfm.models.unet import UNetModel
 

diff --git a/tests/test_optimal_transport.py b/tests/test_optimal_transport.py
@@ -71,17 +71,19 @@ def test_wasserstein(batch_size=128, seed=1980):
     M = torch.cdist(x0.reshape(x0.shape[0], -1), x1.reshape(x1.shape[0], -1))
     pot_W22 = ot.emd2(ot.unif(x0.shape[0]), ot.unif(x1.shape[0]), (M**2).numpy())
     pot_W2 = np.sqrt(pot_W22)
-    W2 = wasserstein(x0, x1, 'exact')
+    W2 = wasserstein(x0, x1, "exact")
 
     pot_W1 = ot.emd2(ot.unif(x0.shape[0]), ot.unif(x1.shape[0]), M.numpy())
-    W1 = wasserstein(x0, x1, 'exact', power=1)
+    W1 = wasserstein(x0, x1, "exact", power=1)
 
-    pot_eot = ot.sinkhorn2(ot.unif(x0.shape[0]), ot.unif(x1.shape[0]), M.numpy(), reg=0.01, numItermax=int(1e7))
-    eot = wasserstein(x0, x1, 'sinkhorn', reg=0.01, power=1)
+    pot_eot = ot.sinkhorn2(
+        ot.unif(x0.shape[0]), ot.unif(x1.shape[0]), M.numpy(), reg=0.01, numItermax=int(1e7)
+    )
+    eot = wasserstein(x0, x1, "sinkhorn", reg=0.01, power=1)
 
     with pytest.raises(ValueError) as excinfo:
-        eot = wasserstein(x0, x1, 'noname', reg=0.01, power=1)
+        eot = wasserstein(x0, x1, "noname", reg=0.01, power=1)
 
-    assert pot_W2==W2
-    assert pot_W1==W1
-    assert pot_eot==eot
+    assert pot_W2 == W2
+    assert pot_W1 == W1
+    assert pot_eot == eot