title: Alternatively train multi task learning model in pytorch - weight updating
body: <p>I want to build a multi task learning model on two related datasets with different inputs and targets. The two tasks are sharing lower-level layers but with different header layers, a minimal example:</p>
  <pre class="lang-py prettyprint-override"><code>class MultiMLP(nn.Module):
      &quot;&quot;&quot;
      A simple dense network for MTL on hard parameter sharing.
      &quot;&quot;&quot;
      def __init__(self):
          super().__init__()
          self.hidden = nn.Linear(100, 200)
          self.out_task0= nn.Linear(200, 1)
          self.out_task0= nn.Linear(200, 1)
  
      def forward(self, x):
          x = self.hidden(x)
          x = F.relu(x)
          y_task0 = self.out_task0(x)
          y_task1 = self.out_task1(x)
          return [y_task0, y_task1]
  </code></pre>
  <p>The dataloader is constructed so that the batches are alternatively generated from two datasets, i.e. <em>batch 0, 2, 4, ...</em> from <em>task 0</em>, <em>batch 1, 3, 5, ...</em> from <em>task 1</em>.  I wanted to train the network in this way: only update weights for <code>hidden</code> layer and <code>out_task0</code> for batches from <em>task 0</em>, and update only <code>hidden</code> and <code>out_task1</code> for <em>task 1</em>.</p>
  <p>I then alternatively switch <code>requires_grad</code> for the corresponding tasks during training as following. <strong>But I observed that all weights are updated for every iteration.</strong></p>
  <pre class="lang-py prettyprint-override"><code>...
  criterion = MSELoss()
  for i, data in enumerate(combined_loader):
      x, y = data[0], data[1]    
      optimizer.zero_grad()
      # controller is 0 for task0, 1 for task1
      # altenate the header layer
      controller = i % 2
      task0_mode = True if controller == 0 else False
      for name, param in model.named_parameters():
          if name in ['out_task0.weight', 'out_task0.bias']:
              param.requires_grad = task0_mode
          elif name in ['out_task1.weight', 'out_task1.bias']:
              param.requires_grad = not task0_mode
  
      outputs = model(x)[controller]
      loss = criterion(outputs, y)
      loss.backward()
      optimizer.step()
      
      # Monitor the parameter updates
      for name, p in model.named_parameters():
          if name in ['out_task0.weight', 'out_task1.weight']:
              print(f&quot;Controller: {controller}&quot;)
              print(name, p)
  </code></pre>
  <p>Did I miss anything in the training procedure? Or the overall setup will not work?</p>