try fix

FluxML · May 28, 2024 · eb62e15 · eb62e15
1 parent e17bef8
commit eb62e15
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/train.jl b/src/train.jl
@@ -121,7 +121,7 @@ _applyloss(loss, model, d...) = loss(model, d...)
 
 Like [`train!](@ref), but gradient computed in place using [Enzyme](github.com/EnzymeAD/Enzyme.jl)        
 """
-function train!(loss, model_and_shadow::Enzyme.Duplicated, data, opt_state)
+function train!(loss, model_and_shadow::Enzyme.Duplicated, data, opt_state::T) where T
   @withprogress for (i,d) in enumerate(data)
     d_splat = d isa Tuple ? d : (d,)
     _make_zero!(model_and_shadow.dval)