Profiler tf native training

smdebug/core/hook.py

@@ -558,18 +558,19 @@ def _cleanup(self):
         if self.first_process is True:
             remove_claim_file(self.out_dir)
 
-    def _increment_step(self):
+    def _increment_step(self, write_state=True):
         # Update the last_state to the last step number that was saved or seen
-        self._write_state()
+        if write_state:
+            self._write_state()
+            self.written_tensor_name_for_step.clear()
+            self._collections_to_save_for_step = None
 
         self.step += 1
         self.mode_steps[self.mode] += 1
-        self.written_tensor_name_for_step.clear()
 
         # Increment Global step number irrespective of what mode it is
         if self.mode != ModeKeys.GLOBAL:
             self.mode_steps[ModeKeys.GLOBAL] = self.step
-        self._collections_to_save_for_step = None
 
     # Called in the internal AWS codebase to determine
     # if a particular tensor value should be saved