docs: def remove_grad_hooks(module, input):
      """ Remove gradient hooks to all of the parameters and monitored vars """
      for hook in list(module.param_hooks.keys()):
          module.param_hooks[hook].remove()
          module.param_hooks.pop(hook)
      for hook in list(module.var_hooks.keys()):
          module.var_hooks[hook].remove()
          module.var_hooks.pop(hook)