code: def add_grad_hooks(net, match_fn=None):
      """Add backward hooks to all parameters to record gradients
  
      Parameters
      ----------
      net : skorch.NeuralNet
        The neural net instance (has to be initialized).
  
      match_fn : callable or None (default=None)
        If not ``None``, this should be a callable/function that takes the name of
        a parameter as input and returns a bool as output, where ``False``
        indicates that this parameter should be excluded.
  
      Returns
      -------
      recs_grad : dict of list of dict
        Data structure containing the recorded gradients. For each module, for each
        training step, for each learnable parameter, there is an entry of the
        recorded gradients.
  
      recs : dict of list of dict
        Data structure containing the recorded parameter updates. For each module,
        for each training step, for each learnable parameter, there is an entry of
        the recorded parameter updates.
  
      handles : list of torch.utils.hooks.RemovableHandle
        These handles are returned by torch when adding hooks. They can be used to
        remove the hooks.
  
      """
      recs_grad = {}
      recs_param_update = {}
      handles = []
  
      for module_name, module in named_modules(net):
          recs_grad[module_name] = []
          recs_param_update[module_name] = []
  
          for param_name, tensor in module.named_parameters():
              if match_fn and not match_fn(param_name):
                  continue
  
              if not tensor.requires_grad:
                  continue
  
              handle = tensor.register_hook(
                  partial(
                      _add_grad_hook,
                      rec_grad=recs_grad[module_name],
                      rec_param_update=recs_param_update[module_name],
                      param_name=param_name,
                      tensor=tensor,
                  )
              )
              handles.append(handle)
  
      return recs_grad, recs_param_update, handles
  
imports: from collections.abc import Mapping
  from functools import partial
  import numpy
  from sklearn.exceptions import NotFittedError
  from sklearn.utils.validation import check_is_fitted
  import torch
  from skorch.exceptions import NotInitializedError
  from skorch.utils import to_numpy
label_desc: Use image node_id 5 for calling a global function with example usage: add_grad_hooks(net, match_fn) and returns: recs_grad,  recs_param_update,  handles