better itr_n

impICNF · Aug 4, 2023 · 6e5a154 · 6e5a154
1 parent cd716a3
commit 6e5a154
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/src/core.jl b/src/core.jl
@@ -1,11 +1,15 @@
 # SciML interface
 
-function callback_f(ps, l, icnf::AbstractFlows, prgr::Progress, itr_n::AbstractVector)
+function callback_f(ps, l, icnf::AbstractFlows, prgr::Progress, itr_n::AbstractArray)
     ProgressMeter.next!(
         prgr;
-        showvalues = [(:loss_value, l), (:iteration, itr_n[]), (:last_update, Dates.now())],
+        showvalues = [
+            (:loss_value, l),
+            (:iteration, only(itr_n)),
+            (:last_update, Dates.now()),
+        ],
     )
-    itr_n[] += one(itr_n[])
+    itr_n[] += one(only(itr_n))
     false
 end
 

diff --git a/src/core_cond_icnf.jl b/src/core_cond_icnf.jl
@@ -97,7 +97,7 @@ function MLJModelInterface.fit(model::CondICNFModel, verbosity, XY)
                     desc = "Fitting (epoch: $ep of $(model.n_epochs)): ",
                     showspeed = true,
                 )
-                itr_n = [1]
+                itr_n = ones(Int)
                 tst_one = @timed res = solve(
                     optprob_re,
                     opt,

diff --git a/src/core_icnf.jl b/src/core_icnf.jl
@@ -95,7 +95,7 @@ function MLJModelInterface.fit(model::ICNFModel, verbosity, X)
                     desc = "Fitting (epoch: $ep of $(model.n_epochs)): ",
                     showspeed = true,
                 )
-                itr_n = [1]
+                itr_n = ones(Int)
                 tst_one = @timed res = solve(
                     optprob_re,
                     opt,