Function create_k_shot_mask

code/utils.py:436–446

(labels, args)

    return loss

def create_k_shot_mask(labels, args):
    k = args.shot
    unique_classes = torch.unique(labels)
    k_shot_mask = torch.zeros_like(labels, dtype=torch.bool)

    for cls in unique_classes:
        class_indices = (labels == cls).nonzero(as_tuple=True)[0]
        chosen_indices = class_indices[torch.randperm(len(class_indices))[:k]]
        k_shot_mask[chosen_indices] = True

    return k_shot_mask