@@ -46,6 +46,7 @@ function train_policy!(
4646 epochs= 100 ,
4747 metrics:: Tuple = (),
4848 maximizer_kwargs= sample -> sample. context,
49+ verbose= false ,
4950)
5051 (; nb_samples, ε, threaded, training_optimizer, seed) = algorithm
5152 (; statistical_model, maximizer) = policy
@@ -70,7 +71,7 @@ function train_policy!(
7071 push! (history, :training_loss , 0 , evaluate! (train_loss_metric, context))
7172 evaluate_metrics! (history, metrics, context)
7273
73- @showprogress for epoch in 1 : epochs
74+ @showprogress enabled = verbose for epoch in 1 : epochs
7475 next_epoch! (context)
7576 for batch in train_dataset
7677 val, grads = Flux. withgradient (statistical_model) do m
@@ -107,6 +108,7 @@ function train_policy!(
107108 epochs= 100 ,
108109 metrics:: Tuple = (),
109110 maximizer_kwargs= sample -> sample. context,
111+ verbose= false ,
110112)
111113 data_loader = DataLoader (train_dataset; batchsize= 1 , shuffle= false )
112114 return train_policy! (
@@ -116,5 +118,6 @@ function train_policy!(
116118 epochs= epochs,
117119 metrics= metrics,
118120 maximizer_kwargs= maximizer_kwargs,
121+ verbose= verbose,
119122 )
120123end
0 commit comments