setup only on AbstractRules

LuxDL · Oct 14, 2022 · fe19182 · fe19182
1 parent 9a34140
commit fe19182
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "0.4.28"
+version = "0.4.29"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/src/utils.jl b/src/utils.jl
@@ -143,7 +143,7 @@ function Functors.functor(::Type{<:ComponentArray}, c)
            ComponentArray
 end
 
-Optimisers.setup(opt, ps::ComponentArray) = Optimisers.setup(opt, getdata(ps))
+Optimisers.setup(opt::AbstractRule, ps::ComponentArray) = Optimisers.setup(opt, getdata(ps))
 
 function Optimisers.update(tree, ps::ComponentArray, gs::ComponentArray)
     tree, ps_new = Optimisers.update(tree, getdata(ps), getdata(gs))