How update weights of two separate neural network with a computed loss?
I have an encoder and a proxy network that help the encoder to maximize information between its input(an image) and output (feature vector of image). to get this done, I used a loss function that estimate MI and by an optimizer the weights of both networks get updated with computed loss, but I'm not sure that does this get done correctly or not. I used following code (in pytorch):

# Clear the previous gradients
discriminator_net_optim.zero_grad()
encoder_net_optim.zero_grad()

autograd.backward(loss)

torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 2)
torch.nn.utils.clip_grad_norm_(encoder.parameters(), 2)

# adjust weights in discriminator and encoder
discriminator_net_optim.step()
encoder_net_optim.step()


any help or suggestion is appreciated.

If you have multiple networks, this is an example of how they would train

encoder = Encoder(args).to(device)
decoder = Decoder(args).to(device)

params = list(encoder.parameters()) + list(decoder.parameters())
optimizer = torch.optim.Adam(params, learning_rate)


And this is called on each batch:

 optimizer.zero_grad()
 loss.backward()
 optimizer.step()


Hope it helps
