experiments
/
STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

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
This commit is contained in:
Ran Chen 2020-06-12 12:18:51 -07:00 committed by TensorFlower Gardener
parent 2bc49eade6
commit 5ec77d721d
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/python/distribute/multi_process_lib.py
View File

@ -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()