fix prof switch

pull/1091/head
Hui Zhang 4 years ago
parent 3e8f802e4e
commit 1a64d350c6

@ -200,10 +200,8 @@ class Trainer():
batch_sampler.set_epoch(self.epoch)
def after_train_batch(self):
if self.args.profiler_options:
profiler.add_profiler_step(self.args.profiler_options)
if self.args.benchmark_max_step and self.iteration > self.args.benchmark_max_step:
profiler.add_profiler_step(self.args.profiler_options)
logger.info(
f"Reach benchmark-max-step: {self.args.benchmark_max_step}")
sys.exit(

Loading…
Cancel
Save