def body(self, features):
    """Body of Residual Shuffle-Exchange network.

    Args:
      features: dictionary of inputs and targets

    Returns:
      the network output.
    """

    inputs = tf.squeeze(features["inputs"], axis=2)
    logits = residual_shuffle_network(inputs, self._hparams)
    return tf.expand_dims(logits, axis=2)
Body of Residual Shuffle-Exchange network.

    Args:
      features: dictionary of inputs and targets

    Returns:
      the network output.
    
tensor2tensor/models/research/residual_shuffle_exchange.py