Skip to content

Commit b8978b5

Browse files
committed
bump device memory limit when running tests locally on single gpu
1 parent bc6d9cb commit b8978b5

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

ci/envs/run_tests/bazel_gpu_local

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,3 +29,6 @@ export JAX_ENABLE_X64=0
2929

3030
# Clone XLA at HEAD.
3131
export JAXCI_CLONE_MAIN_XLA=1
32+
33+
# Set per device memory limit to 20.5 GB
34+
export TF_PER_DEVICE_MEMORY_LIMIT_MB=20480

0 commit comments

Comments
 (0)