typos

FluxML · Nov 5, 2024 · 557ad49 · 557ad49
1 parent 6310548
commit 557ad49
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/ext/FluxEnzymeExt/FluxEnzymeExt.jl b/ext/FluxEnzymeExt/FluxEnzymeExt.jl
@@ -1,7 +1,8 @@
 module FluxEnzymeExt
 
 using Flux
-import Flux.Train: _make_zero!, _enzyme_train!, _rule_to_state
+using Flux: _make_zero!
+import Flux.Train: _enzyme_train!, _rule_to_state
 import Flux.Optimise
 import Optimisers
 import Enzyme

diff --git a/src/train.jl b/src/train.jl
@@ -53,7 +53,7 @@ function setup(rule::Optimisers.AbstractRule, model)
     state
 end
 
-setup(rule, model::Duplicated) = setup(rule, model.val)
+setup(rule::Optimisers.AbstractRule, model::Duplicated) = setup(rule, model.val)
 
 """
     train!(loss, model, data, opt_state)
-Original file line number
+Diff line change
@@ Expand Up / @@ -53,7 +53,7 @@ function setup(rule::Optimisers.AbstractRule, model) @@
         state
     end
-    setup(rule, model::Duplicated) = setup(rule, model.val)
+    setup(rule::Optimisers.AbstractRule, model::Duplicated) = setup(rule, model.val)
     """
         train!(loss, model, data, opt_state)
@@ Expand Down @@