import torch


def mask(out_caps, target):
    """ Function that masks the output of the last capsule layer before feeding it to the decoder.

    The function returns only the capsule corresponding to the correct label and discards the others.

    Args:
        out_caps: output of the last capsule layer    Tensor [batch_size, num_classes, dim_capsules]
        target: labels of the dataset                 Tensor [batch_size]
    Returns:
        masked_out: masked output of the last capsule layer   Tensor [batch_size, dim_capsules]
    """

    batch_index = torch.arange(out_caps.size(0))  # indices along batch dimension
    masked_out = out_caps[batch_index, target, :]  # mask with the target

    return masked_out