Merge pull request FluxML#1944 from darsnack/darsnack/progress-fix

kerem-z · Apr 14, 2022 · f97688e · f97688e
2 parents 24f40b6 + f5ecde1
commit f97688e
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/optimise/train.jl b/src/optimise/train.jl
@@ -112,7 +112,8 @@ Multiple callbacks can be passed to `cb` as array.
 """
 function train!(loss, ps::Params, data, opt::AbstractOptimiser; cb = () -> ())
   cb = runall(cb)
-  n = (Base.IteratorSize(typeof(data)) == Base.HasLength()) ? length(data) : 0
+  itrsz = Base.IteratorSize(typeof(data))
+  n = (itrsz == Base.HasLength()) || (itrsz == Base.HasShape{1}()) ? length(data) : 0
   @withprogress for (i, d) in enumerate(data)
     try
       gs = gradient(ps) do
@@ -129,7 +130,7 @@ function train!(loss, ps::Params, data, opt::AbstractOptimiser; cb = () -> ())
         rethrow(ex)
       end
     end
-    @logprogress i / n
+    @logprogress iszero(n) ? nothing : i / n
   end
 end