fix the train function, the output of loss is not itself

DEEPDIP-project · Sep 5, 2024 · 665ad2f · 665ad2f
1 parent 6bfb21b
commit 665ad2f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/simulations/NavierStokes_2D/scripts/NeuralClosure+SciML.jl b/simulations/NavierStokes_2D/scripts/NeuralClosure+SciML.jl
@@ -160,10 +160,10 @@ tstate = Lux.Training.TrainState(closure, θ, st, OptimizationOptimisers.Adam(0.
 _, loss, stats, tstate = Lux.Training.single_train_step!(
     Optimization.AutoZygote(), loss_lux_style, train_data, tstate)
 
-tstate = train(closure, θ, st, dataloader, loss_lux_style;
+train(closure, θ, st, dataloader, loss_lux_style;
     nepochs = 100, ad_type = Optimization.AutoZygote(),
     alg = OptimizationOptimisers.Adam(0.1), cpu = true)
-# still a problem with the train function
+
 
 # * A posteriori dataloader
 # indeed the ioarrays are not useful here, what a bummer! We should come up with a format that would be useful for both a-priori and a-posteriori training. 

diff --git a/src/train.jl b/src/train.jl
@@ -44,7 +44,7 @@ function train(dataloaders,
     (; optstate, θ, callbackstate)
 end
 
-function train(model, ps, st, train_dataloader, loss;
+function train(model, ps, st, train_dataloader, loss_function;
         nepochs = 100, ad_type = Optimization.AutoZygote(),
         alg = OptimizationOptimisers.Adam(0.1), cpu::Bool = false, kwargs...)
     dev = cpu ? Lux.cpu_device() : Lux.gpu_device()
@@ -55,7 +55,7 @@ function train(model, ps, st, train_dataloader, loss;
         #y = dev(y)
         data = train_dataloader()
         _, loss, stats, tstate = Lux.Training.single_train_step!(
-            ad_type, loss, data, tstate)
+            ad_type, loss_function, data, tstate)
     end
     loss, tstate
 end