diff --git a/tests/BUILD b/tests/BUILD
index 6ff21ab9f..86e61cc7f 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -810,7 +810,7 @@ jax_test(
     },
     shard_count = {
         "cpu": 40,
-        "gpu": 30,
+        "gpu": 40,
         "tpu": 40,
     },
     tags = ["noasan"],  # Times out