import torch


def accuracy(logits, labels):
    """
    computes accuracy given logits (computed probabilities) & labels (actual values)
    @params:
        - logits: predictions computed from call to model.forward(...) (Tensor)
        - labels: actual values (labels) (Tensor)
    @returns:
        computed accuracy value
           accuracy = (correct_predictions) / len(labels)
    """
    if logits.size()[1] == 1:
        # binary classification case (just 2 classes)
        predicted = torch.round(logits.data).reshape(-1)
    else:
        vals, predicted = torch.max(logits.data, 1)
    
    total_count = labels.size(0)

    y_pred = predicted.long()
    if len(labels.shape) > 1:
        y_true = labels.reshape(-1).long()  # flatten
    else:
        y_true = labels.long()

    correct_predictions = (y_pred == y_true).sum().item()
    # accuracy is the fraction of correct predictions to total_count
    acc = (correct_predictions / total_count)
    return acc