Differentiate the loss with a more efficient AD backend

[albangossard]

For the moment in the gradient computation, the differentiation of the loss uses Zygote. This is far from being efficient as we know Zygote is slow and this requires a lot of allocations. We could try to differentiate with Enzyme (or any other library which is more efficient than Zygote) and quantify the improvement in terms of run time. If this is significant, we could replace Zygote by another alternative.