Loss.backward() updates loss function that wasn't included

Anwarvic: <p>Right now, I’m trying to train a model using a weighted sum of two loss functions (<code>loss1</code> and <code>loss2</code>) of the same type; as shown below:</p>
<pre><code class="lang-auto">alpha = 0.5
for input, labels in data_loader_train:
    out1, out2 = model(input)
    loss1 = criterion(out1, labels[:, 0])
    loss2 = criterion(out2, labels[:, 1])
    loss = alpha * loss1 + (1-alpha) * loss2
    punc_loss.backward()
    optimizer.step()
    optimizer.zero_grad()
</code></pre>
<p>So far, everything works as expected.</p>
<p>But <code>alpha</code> is a tuning parameter. So, when I set <code>alpha=1</code>, I want to only use the first loss function <code>loss1</code>. But when I do that, I find all parameters related to <code>loss2</code> get updated such as <code>self.fc2</code> &amp; <code>self.bn2</code>. Same when I set <code>alpha=0</code>, I find all parameters related to <code>loss1</code> get updated such as <code>self.fc1</code> &amp; <code>self.bn1</code>.</p>
<p>What is wrong? Any help is much appreciated</p>
<hr>
<p>In case needed, the following is the forward function of the <code>model</code> where <code>self.bert</code> is frozen:</p>
<pre><code class="lang-auto">def forward(self, x):
    x = self.bert(x).hidden_states[-1]
    x = x.view(x.shape[0], -1)
    out1 = self.fc1(self.dropout(self.bn1(x)))
    out2 = self.fc2(self.dropout(self.bn2(x)))
    return out1, out2
</code></pre>

ptrblck: <p>I cannot reproduce the issue and the loss weighting works as expected:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 10)
        self.fc2 = nn.Linear(10, 10)

    def forward(self, x):
        out1 = self.fc1(x)
        out2 = self.fc2(x)
        return out1, out2
    
model = MyModel()
x = torch.randn(1, 10)
criterion = nn.MSELoss()
target = torch.randn(1, 10)

# both losses
alpha = 0.5
out1, out2 = model(x)
loss1 = criterion(out1, target)
loss2 = criterion(out2, target)
loss = alpha * loss1 + (1-alpha) * loss2
loss.backward()

for name, param in model.named_parameters():
    print(name, param.grad.abs().sum() if param.grad != None else None)
model.zero_grad()
# fc1.weight tensor(7.6900)
# fc1.bias tensor(1.3559)
# fc2.weight tensor(7.6103)
# fc2.bias tensor(1.3418)

# loss1
alpha = 1.0
out1, out2 = model(x)
loss1 = criterion(out1, target)
loss2 = criterion(out2, target)
loss = alpha * loss1 + (1-alpha) * loss2
loss.backward()

for name, param in model.named_parameters():
    print(name, param.grad.abs().sum() if param.grad != None else None)
model.zero_grad()
# fc1.weight tensor(15.3799)
# fc1.bias tensor(2.7117)
# fc2.weight tensor(0.)
# fc2.bias tensor(0.)

# loss2
alpha = 0.0
out1, out2 = model(x)
loss1 = criterion(out1, target)
loss2 = criterion(out2, target)
loss = alpha * loss1 + (1-alpha) * loss2
loss.backward()

for name, param in model.named_parameters():
    print(name, param.grad.abs().sum() if param.grad != None else None)
model.zero_grad()
# fc1.weight tensor(0.)
# fc1.bias tensor(0.)
# fc2.weight tensor(15.2206)
# fc2.bias tensor(2.6836)
</code></pre>
<p>Could you post a minimal, executable code snippet showing the issue?<br>
If you are checking for parameter updates directly, note that optimizers with internal states (e.g. <code>Adam</code>) might update parameters with a zero gradient if these parameters were previously updated and thus have a running internal state.</p>

Anwarvic: <p>What an amazing answer from an amazing person!</p>
<p>I think the last paragraph in your reply explains my problem. I was re-running the for-loop with different <code>alpha</code> values. Once I redefined the whole model with the new <code>alpha</code>, everything worked as expected!</p>