update 15

run_image_captioning_flax_reduced.py

@@ -1243,16 +1243,16 @@ def main():
                 state, train_metric = p_train_step(state, batch)
                 train_metrics.append(train_metric)
                 train_time += time.time() - batch_start
-
                 time_per_step = train_time / cur_step
-                _train_metric = unreplicate(train_metric)
-                desc = f"Epoch... ({epoch + 1}/{num_epochs} | Step: {cur_step} | Loss: {_train_metric['loss']} | Learning Rate: {_train_metric['learning_rate']} | Time per step: {time_per_step})"
-                epochs.desc = desc
-                epochs.write(desc)
 
                 # log and save info
                 if training_args.logging_steps > 0 and cur_step % training_args.logging_steps == 0:
 
+                    _train_metric = unreplicate(train_metric)
+                    desc = f"Epoch... ({epoch + 1}/{num_epochs} | Step: {cur_step} | Loss: {_train_metric['loss']} | Learning Rate: {_train_metric['learning_rate']} | Time per step: {time_per_step})"
+                    epochs.desc = desc
+                    epochs.write(desc)
+
                     logger.info(desc)
 
                     with open(os.path.join(training_args.output_dir, "log"), "a", encoding="UTF-8") as fp: