Gradient of a loss function : struggling to avoid arrays mutation

[rakeshvar]

@mothmani
Here is a sample example that I wrote to work around mutating arrays. Or how to define a ChainRule for the offending piece of code. Maybe this can help you fix your problem.
https://github.com/rakeshvar/Zygote-Mutating-Arrays-WorkAround.jl