From f2052739c17aab1cf0fda0d1fcbea7288e28cfdd Mon Sep 17 00:00:00 2001
From: Mike J Innes <mike.j.innes@gmail.com>
Date: Tue, 12 Sep 2017 14:11:03 +0100
Subject: [PATCH] tweaks

---
 src/layers/stateless.jl    | 2 +-
 src/optimise/optimisers.jl | 5 ++++-
 src/optimise/train.jl      | 5 ++++-
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/layers/stateless.jl b/src/layers/stateless.jl
index e1fe49b1f3..12523afec3 100644
--- a/src/layers/stateless.jl
+++ b/src/layers/stateless.jl
@@ -3,5 +3,5 @@
 mse(ŷ, y) = sum((ŷ .- y).^2)/length(y)
 # back!(::typeof(mse), Δ, ŷ, y) = Δ .* (ŷ .- y)
 
-logloss(ŷ, y) = -sum(y .* log.(ŷ))
+logloss(ŷ, y) = -sum(y .* log.(ŷ)) / size(y, 2)
 # back!(::typeof(logloss), Δ, ŷ, y) = 0 .- Δ .* y ./ ŷ
diff --git a/src/optimise/optimisers.jl b/src/optimise/optimisers.jl
index 4de504fca8..95b31b98f1 100644
--- a/src/optimise/optimisers.jl
+++ b/src/optimise/optimisers.jl
@@ -1,5 +1,8 @@
 function descent(p::Param, η::Real)
-  () -> p.x .-= p.Δ .* η
+  function ()
+    p.x .-= p.Δ .* η
+    p.Δ .= 0
+  end
 end
 
 function momentum(p::Param, ρ::Real)
diff --git a/src/optimise/train.jl b/src/optimise/train.jl
index 29916b1396..8ad437dba0 100644
--- a/src/optimise/train.jl
+++ b/src/optimise/train.jl
@@ -7,7 +7,10 @@ tocb(fs::AbstractVector) = () -> foreach(call, fs)
 function train!(m, data, opt; cb = () -> ())
   cb = tocb(cb)
   @progress for x in data
-    back!(m(x...))
+    l = m(x...)
+    isinf(l.data[]) && error("Inf")
+    isnan(l.data[]) && error("NaN")
+    back!(l)
     opt()
     cb()
   end