From 727d120401c305f4e09341b28f1930ba075f7c04 Mon Sep 17 00:00:00 2001
From: Adam Paszke <apaszke@google.com>
Date: Mon, 1 Jul 2024 02:57:39 -0700
Subject: [PATCH] Bump up the shard_count for GPU FFT tests

They seem to be timing out with ASAN and no sharding.

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

diff --git a/tests/BUILD b/tests/BUILD
index 6a4fa6539..b5a9b916d 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -144,6 +144,7 @@ jax_test(
     shard_count = {
         "tpu": 20,
         "cpu": 20,
+        "gpu": 10,
     },
 )