def infer(self, features=None, **kwargs):
    """Custom infer method for Shuffle-Exchange network.

    Args:
      features: Dictionary of inputs and targets
      **kwargs: SE network currently doesn't support auto-regressive output

    Returns:
      dict: Dictionary of outputs.
    """

    del kwargs
    targets = features.get("targets")
    infer_targets = features.get("infer_targets")

    if targets is None and infer_targets is not None:
      features["targets"] = infer_targets

    # Run the model
    self.hparams.force_full_predict = True
    with tf.variable_scope(self.name):
      logits, _ = self.model_fn(features)

    assert len(logits.shape) == 5  # [batch, time, 1, 1, vocab]
    logits = tf.squeeze(logits, [2, 3])
    outputs = tf.argmax(logits, axis=2)

    return {"outputs": outputs, "logits": logits, "scores": None}
Custom infer method for Shuffle-Exchange network.

    Args:
      features: Dictionary of inputs and targets
      **kwargs: SE network currently doesn't support auto-regressive output

    Returns:
      dict: Dictionary of outputs.
    
tensor2tensor/models/research/shuffle_network.py