mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Increase sharding count on nn_test and svd_test to avoid ASAN timeouts.

Browse Source 
PiperOrigin-RevId: 527387005
This commit is contained in:
Skye Wanderman-Milne 2023-04-26 15:10:47 -07:00 committed by jax authors
parent 139e1c2f92
commit 67d80c21cb
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tests/BUILD
View File

@ -148,7 +148,7 @@ jax_test(
shard_count = {
"cpu": 10,
"gpu": 10,
"tpu": 10,
"tpu": 20,
},
)
@ -540,6 +540,9 @@ jax_test(
jax_test(
name = "nn_test",
srcs = ["nn_test.py"],
shard_count = {
"tpu": 10,
},
)
jax_test(