Increase test shard_count for shape_poly_test on GPU

PiperOrigin-RevId: 723940915
jax-ml · Feb 6, 2025 · a678396 · a678396
1 parent 026b6c9
commit a678396
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/BUILD b/tests/BUILD
@@ -1518,7 +1518,7 @@ jax_multiplatform_test(
     ],
     shard_count = {
         "cpu": 4,
-        "gpu": 4,
+        "gpu": 6,
         "tpu": 4,
     },
     tags = [