Using the gradient of wrt the input in the loss function

[gdalle]

I think this is related, ping @avikpal

Nested AD with Lux etc

Nested AD (Starting v0.5.38) Starting v0.5.38, Lux automatically captures some common possibilities of AD calls inside loss functions on lux layers and converts them into a faster version to do a JVP over gradient (instead of the default reverse over reverse). In short, this finally handles several complaints we have had over the years of not being able to handle nested AD of Neural Networks efficiently. See Nested Automatic Differentiation | Lux.jl Documentation function loss_function2(model, …