bugfix: big loss (#118)

impICNF · Jan 16, 2023 · 13149f8 · 13149f8
1 parent f2c5776
commit 13149f8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/core_cond_icnf.jl b/src/core_cond_icnf.jl
@@ -253,7 +253,7 @@ function MLJModelInterface.fit(model::CondICNFModel, verbosity, XY)
     y = convert(model.array_type, y)
     data = DataLoader((x, y); batchsize = model.batch_size, shuffle = true, partial = true)
     ncdata = ncycle(data, model.n_epochs)
-    initial_loss_value = model.loss(model.m, x, y)
+    initial_loss_value = model.loss(model.m, first(data)...)
 
     if model.opt_app isa FluxOptApp
         model.optimizer isa Optimisers.AbstractRule ||
@@ -308,7 +308,7 @@ function MLJModelInterface.fit(model::CondICNFModel, verbosity, XY)
         tst = @timed res = solve(optprob, model.optimizer, ncdata; callback = _callback)
         model.m.p .= res.u
     end
-    final_loss_value = model.loss(model.m, x, y)
+    final_loss_value = model.loss(model.m, first(data)...)
     @info(
         "Fitting",
         "elapsed time (seconds)" = tst.time,

diff --git a/src/core_icnf.jl b/src/core_icnf.jl
@@ -236,7 +236,7 @@ function MLJModelInterface.fit(model::ICNFModel, verbosity, X)
     x = convert(model.array_type, x)
     data = DataLoader((x,); batchsize = model.batch_size, shuffle = true, partial = true)
     ncdata = ncycle(data, model.n_epochs)
-    initial_loss_value = model.loss(model.m, x)
+    initial_loss_value = model.loss(model.m, first(data)...)
 
     if model.opt_app isa FluxOptApp
         model.optimizer isa Optimisers.AbstractRule ||
@@ -291,7 +291,7 @@ function MLJModelInterface.fit(model::ICNFModel, verbosity, X)
         tst = @timed res = solve(optprob, model.optimizer, ncdata; callback = _callback)
         model.m.p .= res.u
     end
-    final_loss_value = model.loss(model.m, x)
+    final_loss_value = model.loss(model.m, first(data)...)
     @info(
         "Fitting",
         "elapsed time (seconds)" = tst.time,