Stop training on Inf/NaN loss (#2070)

* stop training on Inf/NaN loss

* add a test

* improve test

* improve test

* Update train.jl

* Update optimise.jl

src/optimise/train.jl

@@ -1,5 +1,5 @@
 using ProgressLogging: @progress, @withprogress, @logprogress
-import Zygote: Params, gradient
+import Zygote: Params, gradient, withgradient
 
 
 """
@@ -105,8 +105,10 @@ The optimiser should be from the `Flux.Optimise` module (see [Optimisers](@ref))
 Different optimisers can be combined using [`Flux.Optimise.Optimiser`](@ref Flux.Optimiser).
 
 This training loop iterates through `data` once.
+It will stop with a `DomainError` if the loss is `NaN` or infinite.
+
 You can use [`@epochs`](@ref) to do this several times, or 
-use for instance `Iterators.repeat` to make a longer `data` iterator.
+use for instance `Itertools.ncycle` to make a longer `data` iterator.
 
 ## Callbacks
 
@@ -126,9 +128,12 @@ function train!(loss, ps::Params, data, opt::AbstractOptimiser; cb = () -> ())
   n = (itrsz == Base.HasLength()) || (itrsz == Base.HasShape{1}()) ? length(data) : 0
   @withprogress for (i, d) in enumerate(data)
     try
-      gs = gradient(ps) do
+      l, gs = withgradient(ps) do
         loss(batchmemaybe(d)...)
       end
+      if !isfinite(l)
+        throw(DomainError("Loss is $l on data item $i, stopping training"))
+      end
       update!(opt, ps, gs)
       cb()
     catch ex

test/optimise.jl

@@ -87,6 +87,18 @@ end
   Flux.train!(loss, Flux.params(r), (r,), Descent())
 end
 
+@testset "Stop on NaN" begin
+  m = Dense(1 => 1)
+  m.weight .= 0
+  CNT = 0
+  @test_throws DomainError Flux.train!(Flux.params(m), 1:100, Descent(0.1)) do i
+    CNT += 1
+    (i == 51 ? NaN32 : 1f0) * sum(m([1.0]))
+  end
+  @test CNT == 51  # stopped early
+  @test m.weight[1] ≈ -5  # did not corrupt weights
+end
+
 @testset "ExpDecay" begin
 
   @testset "Sanity Check" begin