Support for dictionary losses in model_fn_body to be consistent with model_fn_body_sharded. Also updated inline doc.

PiperOrigin-RevId: 164305140

Author: T2T Team

tensor2tensor/utils/t2t_model.py

@@ -469,7 +469,10 @@ def model_fn_body_sharded(self, sharded_features):
           _with_timing(self.model_fn_body, "model_fn_body"),
           datashard_to_features)
       if isinstance(output, tuple):
-        loss = {"extra": tf.reduce_mean(output[1])}
+        if isinstance(output[1], dict):
+          loss = output[1]
+        else:
+          loss = {"extra": tf.reduce_mean(output[1])}
         output = output[0]
       else:
         loss = {"extra": 0.0}
@@ -483,10 +486,12 @@ def model_fn_body(self, features):
 
     Args:
       features: A dictionary of key to Tensor.  Each Tensor has shape
-         `[batch_size, ?, ?, hidden_size]`.
+         [batch_size, ?, ?, hidden_size].
 
     Returns:
-      a `Tensor` of logits with shape `[batch_size, O, P, body_output_size]`.
+      output: tensor of logits with shape [batch_size, O, P, body_output_size.
+      losses: either single loss as a scalar, a list, a tensor (to be averaged)
+              or a dictionary of losses.
     """
     raise NotImplementedError("Abstract Method")