import torch


def accuracy(output, target):
    """
    Compute accuracy comparing the output of the model and the target.

    Args:
        output  : [batch_size, num_classes, caps_dim] The output from the last caps layer.
        target  : [batch_size] Labels for dataset.

    Returns:
        accuracy (float): The accuracy for a batch.
    """
    batch_size = target.size(0)

    # Compute the norm of the vector capsules
    v_length = torch.sqrt((output ** 2).sum(dim=2))

    # Find the index of the longest vector
    _, max_index = v_length.max(dim=1)

    # vector with 1 where the model makes a correct prediction, 0 where false
    correct_pred = torch.eq(target.cpu(), max_index.data.cpu())

    acc = correct_pred.float().mean()  # mean accuracy of a batch

    return acc