def num_correct_fun(preds, labels):
    """Computes the number of correct predictions."""
    assert preds.size(0) == labels.size(
        0
    ), "Batch dim of predictions and labels must match"
    # Find number of correct predictions
    num_correct = (preds == labels).float().sum()
    return num_correct