Zygote: treating a model output as a constant

Specific Domains Machine Learning
1

Say I have a model m = Dense(1,1), a function f(x) = x, and a loss(a,b,c) = m(a) - b - f(c).

I don’t want zygote to differentiate on f(c) when calling gradient(() -> loss(a,b,c), m) . How do I tell it that ? It is a problem because the function f is generating random numbers and I get weird errors like Abstract type AbstractFloat does not have a definite size..

2

This may also be related to the fact that I’m working with CuArrays

3

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