Internal change

PiperOrigin-RevId: 358944046

official/nlp/transformer/translate.py

@@ -151,14 +151,8 @@ def translate_file(model,
       text = distribution_strategy.run(text_as_per_replica)
       outputs = distribution_strategy.experimental_local_results(
           predict_step(text))
-      tags, unordered_val_outputs = outputs[0]
-      tags = [tag.numpy() for tag in tags._values]
-      unordered_val_outputs = [
-          val_output.numpy() for val_output in unordered_val_outputs._values]
-      # pylint: enable=protected-access
-      val_outputs = [None] * len(tags)
-      for k in range(len(tags)):
-        val_outputs[tags[k]] = unordered_val_outputs[k]
+      val_outputs = [output for _, output in outputs]
+
       val_outputs = np.reshape(val_outputs, [params["decode_batch_size"], -1])
     else:
       val_outputs, _ = model.predict(text)