pass verbosity to collate in update

FluxML · Sep 9, 2024 · 56838bf · 56838bf
1 parent ff32ca5
commit 56838bf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/mlj_model_interface.jl b/src/mlj_model_interface.jl
@@ -175,7 +175,7 @@ function MLJModelInterface.update(model::MLJFluxModel,
         rng = true_rng(model)
         chain = build(model, rng, shape) |> move
         # reset `optimiser_state`:
-        data = move.(collate(model, X, y))
+        data = move.(collate(model, X, y, verbosity))
         nbatches = length(data[2])
         regularized_optimiser = MLJFlux.regularized_optimiser(model, nbatches)
         optimiser_state = Optimisers.setup(regularized_optimiser, chain)