code: def register_on_step_begin(model):
      """Mount on_step_begin to the model.
      Args:
          model (`torch.nn.module`)
              The torch module object to be pruned.
      """
  
      def hook(module, input):
          for pruner in module.pruners:
              pruner.on_step_begin(0)
  
      hook_handle = model.register_forward_pre_hook(hook)
      return hook_handle
  
file_path: DeepSpeed/deepspeed/compression/helper.py
imports: import torch
  from .basic_layer import Embedding_Compress, LinearLayer_Compress, Conv2dLayer_Compress, BNLayer_Compress, ColumnParallelLinear_Compress, RowParallelLinear_Compress
  from .constants import *
  from deepspeed.utils import logger
label_desc: Use image node_id 9 for calling a global function with example usage: register_on_step_begin(model) and returns: hook_handle