I don’t know if the only way to do it would be to, for example, make a new class GANOptimizer where the step() method requires two closures, one for each loss function. That’s pretty messy but I’m not sure of another solution