diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 114cae18bf..5a1bbef01d 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -658,7 +658,7 @@ def run(self) -> None: writer = SummaryWriter(log_dir=self.tensorboard_log_dir) if self.enable_profiler or self.profiling: prof = torch.profiler.profile( - schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1), + schedule=torch.profiler.schedule(wait=1, warmup=15, active=3, repeat=1), on_trace_ready=torch.profiler.tensorboard_trace_handler( self.tensorboard_log_dir ) @@ -1032,10 +1032,14 @@ def log_loss_valid(_task_key="Default"): writer.close() if self.enable_profiler or self.profiling: prof.stop() + if self.enable_profiler: + log.info( + f"The profiling trace has been saved under {self.tensorboard_log_dir}" + ) if self.profiling: prof.export_chrome_trace(self.profiling_file) log.info( - f"The profiling trace have been saved to: {self.profiling_file}" + f"The profiling trace has been saved to: {self.profiling_file}" ) def save_model(self, save_path, lr=0.0, step=0) -> None: