    def get_sample_sizes(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
task_ids: (bsz,)
get sample sizes for each task
examples/attention_head_selection/src/speech_to_text_head_selection.py