add callbacks back

2017-09-01 23:59:44 -04:00 · 2017-09-01 23:59:44 -04:00 · fe2b35facc
commit fe2b35facc
parent bf098d551c
1 changed files with 3 additions and 4 deletions
--- a/src/optimise/train.jl
+++ b/src/optimise/train.jl
@ -1,13 +1,12 @@
 using Flux.Tracker: back!

-function train!(m, data, opt; epoch = 1)
+function train!(m, data, opt; epoch = 1, cb = () -> ())
  for e in 1:epoch
    epoch > 1 && info("Epoch $e")
    for (x, y) in data
-      loss = m(x, y)
-      @show loss
-      back!(loss)
+      back!(m(x, y))
      opt()
+      cb()
    end
  end
 end