Fix the ValueError: Error when checking model input on the new codepath (#7382)

* Fix the ValueError: Error when checking model input on the new codepath

Fixes the following error:

  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training.py", line 2428, in _standardize_user_data
    exception_prefix='input')
  File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training_utils.py", line 530, in standardize_input_data
    str(len(data)) + ' arrays: ' + str(data)[:200] + '...')
ValueError: Error when checking model input: the list of Numpy arrays that you are passing to your model is not the size the model expected. Expected to see 2 array(s), but instead got the following list of 1 arrays: [<tf.Tensor 'cond_8/Identity:0' shape=(None, None) dtype=int64>]..

Tested and reproduced by running trasformer_main_test (thanks to whoever wrote it, phew!)

* Remove the unnecessary TODO.

official/transformer/v2/transformer_main.py

@@ -152,13 +152,10 @@ class TransformerTask(object):
 
     model.summary()
 
-    # TODO(guptapriya): Figure out a way to structure input that works in both
-    # distributed and non distributed cases.
     train_ds = data_pipeline.train_input_fn(params)
-    if not self.distribution_strategy:
-      map_data_fn = data_pipeline.map_data_for_transformer_fn
-      train_ds = train_ds.map(
-          map_data_fn, num_parallel_calls=params["num_parallel_calls"])
+    map_data_fn = data_pipeline.map_data_for_transformer_fn
+    train_ds = train_ds.map(map_data_fn, 
+                            num_parallel_calls=params["num_parallel_calls"])
 
     callbacks = self._create_callbacks(flags_obj.model_dir, 0, params)