def METHOD_NAME(self, enc, pred, log_normalize=False):
    logits = self._model.joint(enc, pred)[:, 0, 0, :]
    if not log_normalize:
        return logits
    probs = F.log_softmax(logits, dim=len(logits.shape) - 1)
    return probs