def METHOD_NAME(self, sample, task_ids, num_tasks):
    """
    task_ids: (bsz,)
    get sample sizes for each task
    """
    bsz = task_ids.size(0)
    mat = torch.zeros((num_tasks, bsz), device=task_ids.device)
    mat[task_ids, torch.arange(bsz)] = 1.0
    ntokens = torch.sum(sample['target'] != 1, dim=-1)
    sample_sizes = torch.matmul(mat, ntokens.float())
    return sample_sizes