Fix bug in test dataset generation

PiperOrigin-RevId: 273066504

Author: rachellj218

official/resnet/ctl/ctl_imagenet_main.py

@@ -128,13 +128,14 @@ def _test_data_fn(ctx=None):
           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