From fa66e731e640acfcafba39cb1acd34e729a06b38 Mon Sep 17 00:00:00 2001
From: Adam Paszke <apaszke@google.com>
Date: Thu, 18 Apr 2024 06:03:40 -0700
Subject: [PATCH] Increase sharding to avoid timeouts

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

diff --git a/tests/BUILD b/tests/BUILD
index 177956d2d..02cce6b31 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -667,6 +667,7 @@ jax_test(
     name = "nn_test",
     srcs = ["nn_test.py"],
     shard_count = {
+        "cpu": 10,
         "tpu": 10,
         "gpu": 10,
     },