fixed step time

Author: NRauschmayr

smdebug/pytorch/hook.py

@@ -495,6 +495,11 @@ def back(grad):
                     self._save_for_tensor(self.GRADIENT_PREFIX + tname, grad)
             self._save_custom_tensors_post_step()
 
+            # update step time
+            now = time.time()
+            if self.step_event:
+                self.step_event.update_end_time(now)
+
         return back
 
     def _backward_apply(self, module):