diff options
Diffstat (limited to 'tensorflow/contrib/lite/testing/generate_examples.py')
-rw-r--r-- | tensorflow/contrib/lite/testing/generate_examples.py | 17 |
1 files changed, 11 insertions, 6 deletions
diff --git a/tensorflow/contrib/lite/testing/generate_examples.py b/tensorflow/contrib/lite/testing/generate_examples.py index 86540d58a6..b122818221 100644 --- a/tensorflow/contrib/lite/testing/generate_examples.py +++ b/tensorflow/contrib/lite/testing/generate_examples.py @@ -36,6 +36,10 @@ import traceback import zipfile import numpy as np from six import StringIO + +# TODO(aselle): Disable GPU for now +os.environ["CUDA_VISIBLE_DEVICES"] = "-1" + import tensorflow as tf from google.protobuf import text_format # TODO(aselle): switch to TensorFlow's resource_loader @@ -379,12 +383,13 @@ def make_zip_of_tests(zip_path, report["toco_log"] = "" tf.reset_default_graph() - try: - inputs, outputs = make_graph(param_dict_real) - except (tf.errors.UnimplementedError, tf.errors.InvalidArgumentError, - ValueError): - report["tf_log"] += traceback.format_exc() - return None, report + with tf.device('/cpu:0'): + try: + inputs, outputs = make_graph(param_dict_real) + except (tf.errors.UnimplementedError, tf.errors.InvalidArgumentError, + ValueError): + report["tf_log"] += traceback.format_exc() + return None, report sess = tf.Session() try: |