Function step

normalizing_flow/main.py:37–41

(x)

    @partial(mx.compile, inputs=state, outputs=state)
    def step(x):
        loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
        loss, grads = loss_and_grad_fn(model, x)
        optimizer.update(model, grads)
        return loss

    with trange(args.n_steps) as steps:
        for it in steps: