import torch
import torch.nn as nn
import torch.optim as optim

class myModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(100,200)
        self.fc2 = nn.Linear(200, 2, bias=False)
        self.fc3 = nn.Linear(200, 3, bias=False)

    def forward(self,x):
        h = self.fc1(x)
        out2 = self.fc2(h)
        out3 = self.fc3(h)
        return out2, out3

    

if __name__ == '__main__':
    torch.autograd.set_detect_anomaly(True)
    some_value = 0.5
    model = myModel()
    optimizer12 = optim.Adam(list(model.fc1.parameters()) + list(model.fc2.parameters()), lr=1e-3, weight_decay=1e-4)
    optimizer3 = optim.Adam(model.fc3.parameters(), lr=1e-3, weight_decay=1e-4)
    criterion2 = nn.CrossEntropyLoss()
    criterion3 = nn.CrossEntropyLoss()
    input = torch.rand((1,100))
    label2 = torch.randint(0,2,(1,))
    label3 = torch.randint(0,3,(1,))
    

    output2, output3 = model(input)
    
    loss2 = criterion2(output2, label2)
    loss3 = criterion3(output3, label3)

    #loss3.backward() # produces  RuntimeError: Trying to backward through the graph a second time
    loss3.backward(retain_graph=True) # produces: RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation
    optimizer3.step() 
    
    loss3 = criterion3(output3, label3)
    loss12 = loss2 + some_value*loss3
    loss12.backward()
    optimizer12.step()