diff --git a/src/mldebug/batch_runner.py b/src/mldebug/batch_runner.py index 2f68124..c14098b 100644 --- a/src/mldebug/batch_runner.py +++ b/src/mldebug/batch_runner.py @@ -422,7 +422,7 @@ def _run_stamp(self, layer, sid, target_itr, cur_it=1): self.state.error = True break - if sid == 0: + if sid == 0 and not self.args.run_flags.skip_iter2: self.dumper.dump_l3_buffers(layer) return not self.state.error