changes as requested

src/Flux.jl

@@ -21,7 +21,7 @@ include("optimise/Optimise.jl")
 using .Optimise
 using .Optimise: @epochs
 export SGD, ADAM, ADAMW, AdaMax, Momentum, Nesterov,
-       RMSProp, ADAGrad, ADADelta, AMSGrad, NADAM, stop, StopException
+       RMSProp, ADAGrad, ADADelta, AMSGrad, NADAM, stop
 
 include("utils.jl")
 include("onehot.jl")

src/optimise/Optimise.jl

@@ -2,7 +2,7 @@ module Optimise
 
 export train!,
   SGD, ADAM, ADAMW, AdaMax, Momentum, Nesterov,
-  RMSProp, ADAGrad, ADADelta, AMSGrad, NADAM, stop, StopException
+  RMSProp, ADAGrad, ADADelta, AMSGrad, NADAM
 
 struct Param{T}
   x::T

src/optimise/train.jl

@@ -55,14 +55,13 @@ function train!(loss, data, opt; cb = () -> ())
   cb = runall(cb)
   opt = runall(opt)
   @progress for d in data
-    l = loss(d...)
-    @interrupts back!(l)
-    opt()
     try
-      cb()
+      l = loss(d...)
+      @interrupts back!(l)
+      opt()
+      cb() == :stop && break
     catch ex
       if ex isa StopException
-        @info "Stop condition met"
         break
       else
         rethrow(ex)