Fixing namespace bug and requirenment

PiperOrigin-RevId: 338221474
google-research · Oct 21, 2020 · 9897901 · 9897901
1 parent cd7893e
commit 9897901
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/rigl/cifar_resnet/data_helper.py b/rigl/cifar_resnet/data_helper.py
@@ -102,7 +102,7 @@ def parse_serialized_example(record):
   # example is seen twice.
   dataset = dataset.batch(batch_size)
 
-  images_batch, labels_batch = tf.compat.v1.data.make_one_shot_iterator(
+  images_batch, labels_batch = tf.data.make_one_shot_iterator(
       dataset).get_next()
 
   return (tf.reshape(images_batch, [batch_size, IMG_SIZE, IMG_SIZE, 3]),

diff --git a/rigl/requirements.txt b/rigl/requirements.txt
@@ -2,4 +2,4 @@ absl-py>=0.6.0
 numpy>=1.15.4
 six>=1.12.0
 tensorflow>=1.12.0,<2.0  # change to 'tensorflow-gpu' for gpu support
-tensorflow-datasets
+tensorflow-datasets==2.1