class_docstr: A model wrapper to run model forward/backward steps and storing some intermediate feature/gradient information.
docstr: Args:
      nn_module: the model to be wrapped.
      target_layer_names: the names of the layer to cache.
      register_forward: whether to cache the forward pass output corresponding to `target_layer_names`.
      register_backward: whether to cache the backward pass output corresponding to `target_layer_names`.
code: def __init__(
      self,
      nn_module: nn.Module,
      target_layer_names: str | Sequence[str],
      register_forward: bool = False,
      register_backward: bool = False,
  ):
      """
  
      Args:
          nn_module: the model to be wrapped.
          target_layer_names: the names of the layer to cache.
          register_forward: whether to cache the forward pass output corresponding to `target_layer_names`.
          register_backward: whether to cache the backward pass output corresponding to `target_layer_names`.
      """
      self.model = nn_module
      self.target_layers = ensure_tuple(target_layer_names)
  
      self.gradients: dict[str, torch.Tensor] = {}
      self.activations: dict[str, torch.Tensor] = {}
      self.score: torch.Tensor | None = None
      self.class_idx: int | None = None
      self.register_backward = register_backward
      self.register_forward = register_forward
  
      _registered = []
      for name, mod in nn_module.named_modules():
          if name not in self.target_layers:
              continue
          _registered.append(name)
          if self.register_backward:
              if pytorch_after(1, 8):
                  if (
                      "inplace" in mod.__dict__
                      and mod.__dict__["inplace"]
                  ):
                      # inplace=True causes errors for register_full_backward_hook
                      mod.__dict__["inplace"] = False
                  mod.register_full_backward_hook(
                      self.backward_hook(name)
                  )
              else:
                  mod.register_backward_hook(self.backward_hook(name))
          if self.register_forward:
              mod.register_forward_hook(self.forward_hook(name))
      if self.target_layers and (
          len(_registered) != len(self.target_layers)
      ):
          warnings.warn(
              f"Not all target_layers exist in the network module: targets: {self.target_layers}."
          )
  
imports: from __future__ import annotations
  import warnings
  from collections.abc import Callable, Sequence
  from typing import cast
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from monai.config import NdarrayTensor
  from monai.transforms import ScaleIntensity
  from monai.utils import ensure_tuple, pytorch_after
  from monai.visualize.visualizer import default_upsampler
label_desc: Use image node_id 1 to create a new ModelWithHooks object  with example: obj = ModelWithHooks(nn_module, target_layer_names, register_forward, register_backward)