Lux.jl is certainly a good option. Depending on what you need exactly, you may also be fine with just using DifferentiationInterface.jl (for autodiff), any optimization package (Optimisers.jl, Optim.jl, Optimization.jl), and custom loss functions or the losses in Flux.jl, Lux.jl or LossFunctions.jl
1 Like