Merge pull request #2054 from christiangnrd/0.13.5_caching_fix

Make params non-differentiable (Closes #2040 & #2048)
FluxML · Aug 30, 2022 · 31e4dd0 · 31e4dd0
2 parents c04210c + f5793b5
commit 31e4dd0
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/functor.jl b/src/functor.jl
@@ -88,6 +88,9 @@ function params(m...)
   return ps
 end
 
+# Allows caching of the parameters when params is called within gradient() to fix #2040.
+@non_differentiable params(m...)
+
 struct FluxCUDAAdaptor end
 adapt_storage(to::FluxCUDAAdaptor, x) = CUDA.cu(x)
 adapt_storage(to::FluxCUDAAdaptor, x::Zygote.FillArrays.AbstractFill) = CUDA.cu(collect(x))