Pullback for both Flux model's inputs and parameters

One way is to avoid Flux.params in favour of explicit gradients, and write

grad_in, grad_model = gradient(|>, inputs, model)

If you must use implicit gradients, then I believe you could write something like

grad = gradient(() -> model(inputs), Flux.params((model, inputs)))

To be more precise you may need to tell us what inputs are.

(These could both use withgradient or pullback instead.)

3 Likes