import torch


def extract_indices(seq, which_one):
    """
    Helper method to extract the indices of elements which have the specified label.
    :param labels: (torch.tensor) Labels of the context set.
    :param which_class: Label for which indices are extracted.
    :return: (torch.tensor) Indices in the form of a mask that indicate the locations of the specified label.
    """
    mask = torch.eq(seq, which_one)  # binary mask of labels equal to which_class
    mask_indices = torch.nonzero(mask, as_tuple=False)  # indices of labels equal to which class
    return torch.reshape(mask_indices, (-1,))