docstr: task_ids: (bsz,)
  get sample sizes for each task
docstr_tok: ["task_ids",":","(","bsz",",",")","get","sample","sizes","for","each","task"]
code: 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
  
file_path: fairseq/examples/attention_head_selection/src/speech_to_text_head_selection.py
imports: import torch
  from fairseq.optim.amp_optimizer import AMPOptimizer
  from fairseq.tasks import register_task
  from fairseq.tasks.speech_to_text import SpeechToTextTask
  from .data.speech_to_text_dataset_with_domain import SpeechToTextDatasetCreatorWithDomain
  from .loss.attention_head_selection import HeadSelectionLoss
label_desc: Use image node_id 6 for calling the SpeechToTextHeadSelectionTask obj's underlying member method code with example usage: obj.get_sample_sizes(sample, task_ids, num_tasks) and returns: sample_sizes