Revert "Set input layer `batch_size` in multi-replica mode" (#6598)

* Revert "Set input layer `batch_size` in multi-replica mode (#6578)"

This reverts commit f1a59682.

* Rename variables.

official/resnet/keras/keras_imagenet_main.py

@@ -176,7 +176,12 @@ def run(flags_obj):
     if flags_obj.enable_xla and not flags_obj.enable_eager:
       # TODO(b/129861005): Fix OOM issue in eager mode when setting
       # `batch_size` in keras.Input layer.
-      input_layer_batch_size = flags_obj.batch_size
+      if strategy and strategy.num_replicas_in_sync > 1:
+        # TODO(b/129791381): Specify `input_layer_batch_size` value in
+        # DistributionStrategy multi-replica case.
+        input_layer_batch_size = None
+      else:
+        input_layer_batch_size = flags_obj.batch_size
     else:
       input_layer_batch_size = None