Currently the training loop is constructed outside the normalizing flow model.
Since the function make_training_loop only has one argument, which is the optix optimizer, it is should be reasonably easy to incorporate the train_flow function into the normalizing flow model, so in the ideal case, the user can do something like nf_model.fit(data).
Currently the training loop is constructed outside the normalizing flow model.
Since the function
make_training_looponly has one argument, which is the optix optimizer, it is should be reasonably easy to incorporate thetrain_flowfunction into the normalizing flow model, so in the ideal case, the user can do something likenf_model.fit(data).