import torch


def correct(output, target, tokens):
    """Computes the number of correct predictions.

    Args:
        output (torch.tensor): output of the model
        target (torch.tensor): masked labels
        tokens (int): vocabulary size

    Returns:
        int: number of correct predictions
    """
    with torch.no_grad():
        mask = target.type(torch.bool)
        labels = torch.masked_select(target, mask)
        mask = mask.unsqueeze(-1).expand_as(output)
        output = torch.masked_select(output, mask).reshape(-1, tokens)
        _, predicted = torch.max(output, dim=-1)
    return (predicted == labels).sum().item()