code: def _build_task_heads(self, head_modules):
          
          # Make task head layer assignments
          num_layers = len(self.config["layer_out_dims"])
          task_head_layers = self._set_task_head_layers(num_layers)
  
          # task_head_layers stores the layer whose output is input to task head t
          # task_map stores the task heads that appear at each layer
          self.task_map = defaultdict(list)
          for t, l in enumerate(task_head_layers):
              self.task_map[l].append(t)
  
          if any(l == 0 for l in task_head_layers) and head_modules is None:
              raise Exception(
                  "If any task head is being attached to layer 0 "
                  "(the input modules), then you must provide a t-length list of "
                  "head_modules, since the output dimension of each input_module "
                  "cannot be inferred."
              )
  
          # Construct heads
          head_dims = [self.K[t] for t in range(self.t)]
  
          heads = nn.ModuleList()
          for t in range(self.t):
              input_dim = self.config["layer_out_dims"][task_head_layers[t]]
              if self.config["pass_predictions"]:
                  for p in self.task_graph.parents[t]:
                      input_dim += head_dims[p]
              output_dim = head_dims[t]
  
              if head_modules is None:
                  head = nn.Linear(input_dim, output_dim)
              elif isinstance(head_modules, list):
                  head = head_modules[t]
              else:
                  head = copy.deepcopy(head_modules)
              heads.append(head)
          return heads
docs: Creates and attaches task_heads to the appropriate network layers