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

add multiaccelerator tag to test

Browse Source
This commit is contained in:
Cjkkkk 2024-03-15 13:22:45 -07:00
parent 4ac1503bd5
commit 85cbe05f25
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

@ -1437,7 +1437,10 @@ jax_test(
"tpu",
"cpu",
],
shard_count = 4,
shard_count = {
"gpu": 4,
},
tags = ["multiaccelerator"],
deps = [
"//jax:fused_attention_stablehlo",
],