From a678396f44c2afabf55ab220ee3de75f3f73cec1 Mon Sep 17 00:00:00 2001
From: George Necula <necula@google.com>
Date: Thu, 6 Feb 2025 08:13:12 -0800
Subject: [PATCH] Increase test shard_count for shape_poly_test on GPU

PiperOrigin-RevId: 723940915
---
 tests/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/BUILD b/tests/BUILD
index e1c055ed5..07ba4b52f 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -1518,7 +1518,7 @@ jax_multiplatform_test(
     ],
     shard_count = {
         "cpu": 4,
-        "gpu": 4,
+        "gpu": 6,
         "tpu": 4,
     },
     tags = [