'Mixture-of-experts models will override this function.
Compute model body on all datashards.
Args:
sharded_features: map from string to list of Tensors each with shape
[batch, ?, ?, body_input_size]
Returns:
sharded_body_output:
a list of Tensors, each with shape [batch, O, P, body_output_size]
extra_loss: a Scalar.'

def model_fn_body_sharded(self, sharded_features):

    with tf.name_scope('model'):
        datashard_to_features = [{k: v[d] for (k, v) in six.iteritems(sharded_features)} for d in xrange(self._num_datashards)]
        output = self._data_parallelism(_with_timing(self.model_fn_body, 'model_fn_body'), datashard_to_features)
        if isinstance(output, tuple):
            losses_sharded = output[1]
            if isinstance(losses_sharded[0], dict):
                loss = {}
                for k in losses_sharded[0].keys():
                    k_loss_sharded = [losses[k] for losses in losses_sharded]
                    loss[k] = tf.reduce_mean(k_loss_sharded)
            else:
                loss = {'extra': tf.reduce_mean(losses_sharded)}
            output = output[0]
        else:
            loss = {'extra': 0.0}
        return (output, loss)
