From 24a3bce4954f15a8d4c2895f670a35e7b7bb112e Mon Sep 17 00:00:00 2001
From: Dhairya Gandhi <dhairya@juliacopmuting.com>
Date: Fri, 17 Aug 2018 17:46:13 +0530
Subject: [PATCH] added stop to break training loop

---
 src/optimise/train.jl | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/optimise/train.jl b/src/optimise/train.jl
index 95009444..0b54ccfa 100644
--- a/src/optimise/train.jl
+++ b/src/optimise/train.jl
@@ -1,5 +1,6 @@
 using Juno
 using Flux.Tracker: back!
+include("../utls.jl")
 
 runall(f) = f
 runall(fs::AbstractVector) = () -> foreach(call, fs)
@@ -33,7 +34,12 @@ The callback can return `:stop` to interrupt the training loop.
 Multiple optimisers and callbacks can be passed to `opt` and `cb` as arrays.
 """
 function train!(loss, data, opt; cb = () -> ())
-  cb = runall(cb)
+  cb = try:
+        runall(cb)
+      catch e
+        if e isa StopException || rethrow()
+          @info "Stop Condition Met"
+          break
   opt = runall(opt)
   @progress for d in data
     l = loss(d...)