code: def add_model_hooks(model: torch.nn.Module):
      def start_time_hook(module, input):
          if hasattr(module, "stage") and module.stage == "decode":
              return
          elif hasattr(module, "stage") and module.stage == "prefill":
              torch.cuda.synchronize()
              module.__start_time__ = time.time()
  
      def end_time_hook(module, input, output):
          if hasattr(module, "stage") and module.stage == "decode":
              return
          elif hasattr(module, "stage") and module.stage == "prefill":
              torch.cuda.synchronize()
              module.__duration__ = time.time() - module.__start_time__
              module.stage = "decode"
  
      if not hasattr(model, "__start_time_hook_handle"):
          model.__start_time_hook_handle__ = (
              model.register_forward_pre_hook(
                  start_time_hook,
              )
          )
  
      if not hasattr(model, "__end_time_hook_handle__"):
          model.__end_time_hook_handle__ = model.register_forward_hook(
              end_time_hook,
          )
  
file_path: deepspeedexamples/inference/huggingface/zero_inference/utils.py
imports: import torch
  import time
label_desc: Use image node_id 13 for calling a global function with example usage: add_model_hooks(model) without return types