bump device memory limit when running tests locally on single gpu

nitins17 · nitins17 · commit b8978b5b68d4 · 2024-09-25T17:24:38.000Z
diff --git a/ci/envs/run_tests/bazel_gpu_local b/ci/envs/run_tests/bazel_gpu_local
@@ -29,3 +29,6 @@ export JAX_ENABLE_X64=0
 
 # Clone XLA at HEAD.
 export JAXCI_CLONE_MAIN_XLA=1
+
+# Set per device memory limit to 20.5 GB
+export TF_PER_DEVICE_MEMORY_LIMIT_MB=20480
