Switch to update! only

Author: darsnack

src/optimise/Optimise.jl

@@ -6,7 +6,7 @@ import Zygote
 import Zygote: Params, gradient
 using AbstractDifferentiation
 import Optimisers
-import Optimisers: update, update!
+import Optimisers: update!
 using LinearAlgebra
 import ArrayInterface
 using ProgressLogging: @progress, @withprogress, @logprogress

src/optimise/train.jl

@@ -24,7 +24,6 @@ function Optimisers.update!(opt::AbstractOptimiser, xs::Params, gs)
 
   return opt, xs
 end
-Optimisers.update(opt::AbstractOptimiser, xs::Params, gs) = update!(opt, xs, gs)
 
 # Callback niceties
 call(f, xs...) = f(xs...)
@@ -139,7 +138,7 @@ function train!(loss, ad::AD.AbstractBackend, model, data, optstate; cb = () ->
     try
       _loss = _build_loss(ad, loss, batchmemaybe(d))
       gs = _gradient_only(AD.gradient(ad, _loss, model))
-      optstate, model = update(optstate, model, gs)
+      optstate, model = update!(optstate, model, gs)
       cb()
     catch ex
       if ex isa StopException