fix

run_image_captioning_flax.py

@@ -1288,6 +1288,10 @@ def main():
 
             if metrics:
 
+                # Save metrics (only for the evaluation/prediction being done along with training)
+                if has_tensorboard and training_args.do_train:
+                    write_metric(summary_writer, metrics, train_time=None, step=cur_step, metric_key_prefix=metric_key_prefix)
+
                 # save final metrics in json
                 metrics = {f"{metric_key_prefix}_{metric_name}": round(value.item(), 6) for metric_name, value in metrics.items()}
                 _path = os.path.join(training_args.output_dir, ckpt_dir, f"{metric_key_prefix}_results.json")
@@ -1298,10 +1302,6 @@ def main():
                 with open(os.path.join(training_args.output_dir, 'log'), 'a', encoding='UTF-8') as fp:
                     fp.write(desc + '\n')
 
-                # Save metrics (only for the evaluation/prediction being done along with training)
-                if has_tensorboard and training_args.do_train:
-                    write_metric(summary_writer, metrics, train_time=None, step=cur_step, metric_key_prefix=metric_key_prefix)
-
             # Save generations
             if generations:
                 with open(os.path.join(training_args.output_dir, ckpt_dir, f'{metric_key_prefix}_generation.json'), 'w', encoding='UTF-8') as fp: