accuracy.py

def accuracy(self, y_hat, Y):
    if self.library == "tf":
        # calculate argmax
        argmax = tf.cast(tf.argmax(y_hat, axis=1), Y.dtype)
        acc = tf.math.reduce_sum(tf.cast(argmax == Y, tf.int32)) / Y.shape[0]
    else:
        argmax = torch.argmax(y_hat, dim=1)
        acc = torch.sum(torch.eq(argmax, Y)) / Y.shape[0]

    return acc