def predict_proba(self,
                    input_fn=None,
                    batch_size=None):
    """Returns prediction probabilities for given features.

    Args:
      input_fn: Input function.
      batch_size: Override default batch size.

    Returns:
      An iterable of predicted probabilities with shape [batch_size, n_classes].
    """
    key = prediction_key.PredictionKey.PROBABILITIES
    preds = self.predict(
        input_fn=input_fn,
        batch_size=batch_size,
        outputs=[key])
    return (pred[key] for pred in preds)