fix prof switch

pull/846/head
Hui Zhang 4 years ago
parent a40312b691
commit 7775abd727

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

Loading…
Cancel
Save