Update test_vmap_max_memory max mem

ott-jax · Feb 18, 2025 · 0d37d1b · 0d37d1b
1 parent e8735e7
commit 0d37d1b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/utils_test.py b/tests/utils_test.py
@@ -164,7 +164,7 @@ def fn(x: jnp.ndarray) -> jnp.ndarray:
 
     np.testing.assert_array_equal(fn(x), x.sum(1))
 
-  @pytest.mark.limit_memory("15MB")
+  @pytest.mark.limit_memory("20MB")
   def test_vmap_max_memory(self, rng: jax.Array):
     n, m, d = 2 ** 16, 2 ** 11, 3
     rng, rng_data = jax.random.split(rng, 2)