Commit 49b90e86 authored by Toby Boyd's avatar Toby Boyd Committed by guptapriya
Browse files

improved v2 check.

parent 308c7934
...@@ -22,6 +22,7 @@ import time ...@@ -22,6 +22,7 @@ import time
import tensorflow as tf import tensorflow as tf
from tensorflow.core.protobuf import rewriter_config_pb2 from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python import tf2
from tensorflow.python.eager import profiler from tensorflow.python.eager import profiler
...@@ -202,9 +203,7 @@ def set_config_v2(enable_xla=False, ...@@ -202,9 +203,7 @@ def set_config_v2(enable_xla=False,
{'layout_optimizer': False} {'layout_optimizer': False}
) )
def is_v2_0(): def is_v2_0():
"""Returns true if using tf 2.0.""" """Returns true if using tf 2.0."""
if hasattr(tf, 'contrib'): return tf2.enabled()
return False
else:
return True
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment