class_docstr: Base pytorch Module with support for task labels.
  
  Multi-task modules are ``torch.nn.Module`` for multi-task
  scenarios. The ``forward`` method accepts task labels, one for
  each sample in the mini-batch.
  
  By default the ``forward`` method splits the mini-batch by task
  and calls ``forward_single_task``. Subclasses must implement
  ``forward_single_task`` or override `forward. If ``task_labels == None``,
  the output is computed in parallel for each task.
class_docstr_tok: ["Base","pytorch","Module","with","support","for","task","labels",".","Multi-task","modules","are","``","torch.nn.Module","``","for","multi-task","scenarios",".","The","``","forward","``","method","accepts","task","labels",",","one","for","each","sample","in","the","mini-batch",".","By","default","the","``","forward","``","method","splits","the","mini-batch","by","task","and","calls","``","forward_single_task","``",".","Subclasses","must","implement","``","forward_single_task","``","or","override","`","forward",".","If","``","task_labels","==","None","``",",","the","output","is","computed","in","parallel","for","each","task","."]
docstr: compute the output given the input `x` and task labels.
  
  :param x:
  :param task_labels: task labels for each sample. if None, the
      computation will return all the possible outputs as a dictionary
      with task IDs as keys and the output of the corresponding task as
      output.
  :return:
docstr_tok: ["compute","the","output","given","the","input","`","x","`","and","task","labels",".",":","param","x",":",":","param","task_labels",":","task","labels","for","each","sample",".","if","None",",","the","computation","will","return","all","the","possible","outputs","as","a","dictionary","with","task","IDs","as","keys","and","the","output","of","the","corresponding","task","as","output",".",":","return",":"]
code: def forward(
      self, x: torch.Tensor, task_labels: torch.Tensor
  ) -> torch.Tensor:
      """compute the output given the input `x` and task labels.
  
      :param x:
      :param task_labels: task labels for each sample. if None, the
          computation will return all the possible outputs as a dictionary
          with task IDs as keys and the output of the corresponding task as
          output.
      :return:
      """
      if task_labels is None:
          return self.forward_all_tasks(x)
  
      if isinstance(task_labels, int):
          # fast path. mini-batch is single task.
          return self.forward_single_task(x, task_labels)
      else:
          unique_tasks = torch.unique(task_labels)
  
      out = torch.zeros(
          x.shape[0], self.max_class_label, device=x.device
      )
      for task in unique_tasks:
          task_mask = task_labels == task
          x_task = x[task_mask]
          out_task = self.forward_single_task(x_task, task.item())
          assert len(out_task.shape) == 2, (
              "multi-head assumes mini-batches of 2 dimensions "
              "<batch, classes>"
          )
          n_labels_head = out_task.shape[1]
          out[task_mask, :n_labels_head] = out_task
      return out
  
imports: import torch
  from torch.nn import Module
  from typing import Optional
  from avalanche.benchmarks.utils.flat_data import ConstantSequence
  from avalanche.benchmarks.scenarios import CLExperience
label_desc: Use image node_id 5 for calling the MultiTaskModule obj's underlying member method code with example usage: obj.forward(x, task_labels) and returns: out, self, self