Fix test

ott-jax · Sep 13, 2024 · 634f3fe · 634f3fe
1 parent 27b639e
commit 634f3fe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/tools/sliced_test.py b/tests/tools/sliced_test.py
@@ -61,7 +61,7 @@ def test_random_projs(
     n, m, dim, n_proj = 12, 17, 5, 13
     rng1, rng2 = jax.random.split(rng, 2)
     a, x, b, y = gen_data(rng1, n, m, dim)
-    weights = jax.random.uniform(rng2, n_proj)
+    weights = jax.random.uniform(rng2, (n_proj,))
 
     # Test non-negative and returns output as needed.
     cost, out = sliced.sliced_wasserstein(