Bump shard_count for shard_map_test to fix timeouts.

PiperOrigin-RevId: 571109311
google · Oct 5, 2023 · a2b70e3 · a2b70e3
1 parent 4fd6bf4
commit a2b70e3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/BUILD b/tests/BUILD
@@ -1171,9 +1171,9 @@ jax_test(
     name = "shard_map_test",
     srcs = ["shard_map_test.py"],
     shard_count = {
-        "cpu": 30,
+        "cpu": 50,
         "gpu": 10,
-        "tpu": 40,
+        "tpu": 50,
     },
     tags = ["multiaccelerator"],
     deps = [