From 5ec77d721da8453ed04cb7d3b7da56070d242939 Mon Sep 17 00:00:00 2001 From: Ran Chen Date: Fri, 12 Jun 2020 12:18:51 -0700 Subject: [PATCH] Set TF_FORCE_GPU_ALLOW_GROWTH for multi process tests Otherwise GPU tests may OOM since TF graps the whole GPUs by default. PiperOrigin-RevId: 316152876 Change-Id: I401aeac1d42e0585b6c6deaa6de4f779e412a020 --- tensorflow/python/distribute/multi_process_lib.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tensorflow/python/distribute/multi_process_lib.py b/tensorflow/python/distribute/multi_process_lib.py index 12b81db7189..9b7851439b7 100644 --- a/tensorflow/python/distribute/multi_process_lib.py +++ b/tensorflow/python/distribute/multi_process_lib.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function import multiprocessing as _multiprocessing +import os import unittest from tensorflow.python.platform import test @@ -42,6 +43,7 @@ class Process(object): def test_main(): """Main function to be called within `__main__` of a test file.""" + os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true' test.main()