Zygote: treating a model output as a constant

[simeonschaub]

Are you looking for Zygote.dropgrad? There was some discussion about a similar problem in this thread.