From bc65639edf0a50bedb93238e8166919720f51b18 Mon Sep 17 00:00:00 2001 From: William Shen Date: Tue, 24 Jan 2023 16:28:00 -0500 Subject: [PATCH] Fix total train time metric logging (#1284) --- nerfstudio/engine/trainer.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index b8aace10f..f24ea5ccd 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -227,14 +227,18 @@ class Trainer: self.save_checkpoint(step) writer.write_out_storage() - # save checkpoint at the end of training - self.save_checkpoint(step) - CONSOLE.rule() - CONSOLE.print("[bold green]:tada: :tada: :tada: Training Finished :tada: :tada: :tada:", justify="center") - if not self.config.viewer.quit_on_train_completion: - CONSOLE.print("Use ctrl+c to quit", justify="center") - self._always_render(step) + # save checkpoint at the end of training + self.save_checkpoint(step) + + # write out any remaining events (e.g., total train time) + writer.write_out_storage() + + CONSOLE.rule() + CONSOLE.print("[bold green]:tada: :tada: :tada: Training Finished :tada: :tada: :tada:", justify="center") + if not self.config.viewer.quit_on_train_completion: + CONSOLE.print("Use ctrl+c to quit", justify="center") + self._always_render(step) @check_main_thread def _always_render(self, step):