log intervals in callback handler (#943)

Summary: Pull Request resolved: #943

Reviewed By: diego-urgell

Differential Revision: D65622296

fbshipit-source-id: ca9136e42045d12b23534afd900419cc413e8383

Author: JKSenthil

torchtnt/framework/_callback_handler.py

@@ -14,6 +14,7 @@
 from torchtnt.framework.callback import Callback
 from torchtnt.framework.state import State
 from torchtnt.framework.unit import TEvalUnit, TPredictUnit, TTrainUnit
+from torchtnt.utils.event_handlers import log_interval
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -124,6 +125,7 @@ def on_exception(
         for cb in callbacks:
             cb.on_exception(state, unit, exc)
 
+    @log_interval("on_train_start", {"category": "callback_handler"})
     def on_train_start(self, state: State, unit: TTrainUnit) -> None:
         fn_name = "on_train_start"
         callbacks = self._callbacks.get(fn_name, [])
@@ -176,12 +178,14 @@ def on_train_step_end(self, state: State, unit: TTrainUnit) -> None:
         for cb in callbacks:
             cb.on_train_step_end(state, unit)
 
+    @log_interval("on_train_epoch_end", {"category": "callback_handler"})
     def on_train_epoch_end(self, state: State, unit: TTrainUnit) -> None:
         fn_name = "on_train_epoch_end"
         callbacks = self._callbacks.get(fn_name, [])
         for cb in callbacks:
             cb.on_train_epoch_end(state, unit)
 
+    @log_interval("on_train_end", {"category": "callback_handler"})
     def on_train_end(self, state: State, unit: TTrainUnit) -> None:
         fn_name = "on_train_end"
         callbacks = self._callbacks.get(fn_name, [])