From 56a97088378b3a33303a23bc21f99754985694d5 Mon Sep 17 00:00:00 2001 From: Hayden Gunraj Date: Tue, 30 Jun 2020 11:43:32 -0400 Subject: [PATCH] Fix purge step (#563) --- tensorboardX/writer.py | 23 ++++++++--------------- 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/tensorboardX/writer.py b/tensorboardX/writer.py index 38278ebe..a6d6d2e5 100644 --- a/tensorboardX/writer.py +++ b/tensorboardX/writer.py @@ -320,23 +320,16 @@ def _get_file_writer(self): return self.file_writer if self.all_writers is None or self.file_writer is None: - if 'purge_step' in self.kwargs.keys(): - most_recent_step = self.kwargs.pop('purge_step') - self.file_writer = FileWriter(logdir=self.logdir, - max_queue=self._max_queue, - flush_secs=self._flush_secs, - filename_suffix=self._filename_suffix, - **self.kwargs) + self.file_writer = FileWriter(logdir=self.logdir, + max_queue=self._max_queue, + flush_secs=self._flush_secs, + filename_suffix=self._filename_suffix, + **self.kwargs) + if self.purge_step is not None: self.file_writer.add_event( - Event(step=most_recent_step, file_version='brain.Event:2')) + Event(step=self.purge_step, file_version='brain.Event:2')) self.file_writer.add_event( - Event(step=most_recent_step, session_log=SessionLog(status=SessionLog.START))) - else: - self.file_writer = FileWriter(logdir=self.logdir, - max_queue=self._max_queue, - flush_secs=self._flush_secs, - filename_suffix=self._filename_suffix, - **self.kwargs) + Event(step=self.purge_step, session_log=SessionLog(status=SessionLog.START))) self.all_writers = {self.file_writer.get_logdir(): self.file_writer} return self.file_writer