Optimize multiple models at once

[google]

Optimize multiple models at once

[noahtren]

One idea is to define a single loss function that takes both models as an input, and then use two optimizers to optimize either model while holding the other constant, like so:

e_optim, e_loss, e_aux = e_optim.optimize(lambda e: loss_fn(e, decoder))
d_optim, d_loss, d_aux = d_optim.optimize(lambda d: loss_fn(encoder, d))

One problem is that this would essentially double training time.

Just thinking out loud and hoping to see if there's a better way to do this with Flax.

[noahtren]

^^ that didn't work. Realized I can just output the loss directly from my model, so I did that instead.