Hi, I am designing an autoencoder with multiple decoders. And I want to update a selected decoder during the training. For example, below is a toy example of the network that has two decoders and a shared encoder.
class myModel(nn.module):
    def __init__(self):
        super(myModel, self).__init__()
        self.ENC = nn.ModuleList()
        self.ENC.append(nn.Conv1d(1,1,kernel_size=2))
        self.DEC1 = nn.ModuleList()
        self.DEC1.append(nn.Conv1d(1,1,kernel_size=2))
        self.DEC2 = nn.ModuleList()
        self.DEC2.append(nn.Conv1d(1,1,kernel_size=2))
    def forward(self, x, decoder_idx):
        for f in self.ENC:
            x = f(x)
        if idx == 1:
            for f in self.DEC1:
                x = f(x)
        elif idx == 2:
            for f in self.DEC2:
                x = f(x)
        return x

model = myModel()

out = model(x, 1)
optimizer.zero_grad()
loss.backward()
optimizer.step()

out = model(x, 2)
optimizer.zero_grad()
loss.backward()
optimizer.step()

When I testify with out = myModel(x, 1), it only updates parameters of ENC and DEC1 as I expected. However, when I try out = myModel(x, 2) after the previous process, it updates ENC and both DEC1 and DEC2. How can I update them selectively in the continuous process?