softmax.py

def softmax(self, logits):

    if self.library == "tf":
        exp = tf.exp(logits)
        denom = tf.math.reduce_sum(exp, 1, keepdims=True)

    else:
        exp = torch.exp(logits)
        denom = torch.sum(exp, dim=1, keepdim=True)

    return exp/denom