Skip to content

GitLab

Commit 1577ed07 authored by Jiri Simsa's avatar Jiri Simsa Committed by A. Unique TensorFlower
Browse files

Fix TF 2.0 tf.data API usage. 

PiperOrigin-RevId: 268401010
parent 5ceb3acf
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
official/recommendation/data_test.py
View file @ 1577ed07
...@@ -127,7 +127,7 @@ class BaseTest(tf.test.TestCase): ...@@ -127,7 +127,7 @@ class BaseTest(tf.test.TestCase):
# type: (tf.data.Dataset, tf.Graph) -> list # type: (tf.data.Dataset, tf.Graph) -> list
with self.session(graph=g) as sess: with self.session(graph=g) as sess:
with g.as_default(): with g.as_default():
batch = dataset.make_one_shot_iterator().get_next() batch = tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()
output = [] output = []
while True: while True:
try: try:
......
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