apply best param on GradientMPIMode and improve epochs print

Author: Giovanni3A

Project.toml

@@ -1,7 +1,7 @@
 name = "ApplicationDrivenLearning"
 uuid = "0856f1c8-ef17-4e14-9230-2773e47a789e"
 authors = ["Giovanni Amorim", "Joaquim Garcia"]
-version = "0.1.2"
+version = "0.1.3"
 
 [deps]
 BilevelJuMP = "485130c0-026e-11ea-0f1a-6992cd14145c"

src/optimizers/gradient.jl

@@ -72,7 +72,8 @@ function train_with_gradient!(
             # store and print cost
             trace[epoch] = C
             if verbose
-                println("Epoch $epoch | Cost = $(round(C, digits=2))")
+                dtime = time() - start_time
+                println("Epoch $epoch | Time = $(round(dtime, digits=1))s | Cost = $(round(C, digits=2))")
             end
 
             # evaluate if best model

src/optimizers/gradient_mpi.jl

@@ -106,7 +106,8 @@ function train_with_gradient_mpi!(
                 # store and print cost
                 trace[epoch] = curr_C
                 if verbose
-                    println("Epoch $epoch | Cost = $(round(curr_C, digits=2))")
+                    dtime = time() - start_time
+                    println("Epoch $epoch | Time = $(round(dtime, digits=1))s | Cost = $(round(curr_C, digits=2))")
                 end
 
                 # evaluate if best model
@@ -129,6 +130,9 @@ function train_with_gradient_mpi!(
         is_done = true
         MPI.bcast(is_done, MPI.COMM_WORLD)
 
+        # fix best model
+        apply_params(model.forecast, best_θ)
+
     elseif JQM.is_worker_process()
         # continuoslly call pmap until controller is done
         while true