Fix Resnet XLA with multi-GPUs (#6510)

Don't pass `batch_size` to keras.layers.Input in DS multi-replica case. There is currently a bug in Keras side which will cause a batch size incompatible error.

official/resnet/keras/keras_imagenet_main.py

@@ -174,9 +174,10 @@ def run(flags_obj):
           optimizer, loss_scale=flags_core.get_loss_scale(flags_obj))
 
     if flags_obj.enable_xla:
-      if strategy:
-          per_replica_batch_size = (
-              flags_obj.batch_size // strategy.num_replicas_in_sync)
+      if strategy and strategy.num_replicas_in_sync > 1:
+        # TODO(b/129791381): Specify `per_replica_batch_size` value in
+        # DistributionStrategy multi-replica case.
+        per_replica_batch_size = None
       else:
         per_replica_batch_size = flags_obj.batch_size
     else: