Fix bug in test dataset generation

PiperOrigin-RevId: 273066504

Fix bug in test dataset generation
PiperOrigin-RevId: 273066504
7acce35d · Jing Li · A. Unique TensorFlower · f8d9c9b8 · 7acce35d
Commit 7acce35d authored Oct 05, 2019 by Jing Li Committed by A. Unique TensorFlower Oct 05, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 7 additions and 6 deletions

official/resnet/ctl/ctl_imagenet_main.py official/resnet/ctl/ctl_imagenet_main.py +7 -6

No files found.
--- a/official/resnet/ctl/ctl_imagenet_main.py
+++ b/official/resnet/ctl/ctl_imagenet_main.py
@@ -128,13 +128,14 @@ def get_input_dataset(flags_obj, strategy):
          input_context=ctx)
      return test_ds

-  if strategy:
-    if isinstance(strategy, tf.distribute.experimental.TPUStrategy):
-      test_ds = strategy.experimental_distribute_datasets_from_function(_test_data_fn)
+    if strategy:
+      if isinstance(strategy, tf.distribute.experimental.TPUStrategy):
+        test_ds = strategy.experimental_distribute_datasets_from_function(
+            _test_data_fn)
+      else:
+        test_ds = strategy.experimental_distribute_dataset(_test_data_fn())
    else:
-      test_ds = strategy.experimental_distribute_dataset(_test_data_fn())
-  else:
-    test_ds = _test_data_fn()
+      test_ds = _test_data_fn()

  return train_ds, test_ds