Update train.jl

kerem-z · May 16, 2024 · d8f1ad8 · d8f1ad8
1 parent 97a490f
commit d8f1ad8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/train.jl b/src/train.jl
@@ -131,7 +131,7 @@ function train!(loss, model_and_shadow::Enzyme.Duplicated, data, opt_state)
       throw(DomainError(lazy"Loss is $l on data item $i, stopping training"))
     end
     opt_state, model = Optimisers.update!(opt_state, model_and_shadow.val, model_and_shadow.dval)
-    model_and_shadow = Duplicated(model, model_and_shadow.dval)
+    model_and_shadow = Enzyme.Duplicated(model, model_and_shadow.dval)
     @logprogress Base.haslength(data) ? i/length(data) : nothing
   end
 end