import torch


def calc_accuracy(logits, labels):
    """
    func to compute accuracy from input logits and labels, meant to be used
    with tensor inputs from Pytorch. This func essentially just implements
    np.sum(np.argmax(logits, axis=0) == labels)/len(labels).
    Make sure if logits is in cuda mode labels need to be in cuda mode.

    Inputs:
        logits tensor - logits of dim num classes x N
        labels tensor - associated labels vector size N
    Outputs:
        train_acc float - float of accuracy between [0, 1]
    
    Example Useage:
        >>> inputs, labels = data # get inputs and labels for current batch
        >>> logits = neural_network(inputs) # forward pass over inputs
        >>> calc_accuracy(logits, labels.cuda())
    """
    max_vals, max_indices = torch.max(logits, 1)
    train_acc = (max_indices == labels).sum().item()/max_indices.size()[0]
    return train_acc