Error in Flux when calculating the gradient due to an in-place operation

I think we mostly need to know what generate_ak does