Problem when training CycleGAN

Sometimes (e.g. CycleGAN) we need to optimize two (or more) models' parameters together because it will be more efficient (e.g. when optimizing cycle loss we definitely don't want to use retain_graph=True).

I was just wondering whether this is the right way to initialize optimizer which aims to optimize both net_a2b's & net_b2a's parameters:

original_optimizer = ...
net_a2b, optimizer, _, _ = deepspeed.initialize(args, net_a2b, original_optimizer)
net_b2a, _, _, _ = deepspeed.initialize(args, net_b2a, original_optimizer)

Any help would be very grateful, thanks in advance!

Start with the two deepspeed.initialize calls in the issue and trace how the optimizer and model parameters are registered. Check the relevant DeepSpeed initialization documentation or entry point, then validate the supported approach with the described CycleGAN case. Done means the supported way to optimize both models together is clearly established without relying on retain_graph=True.