code: def module_rollback_ctx(
      module: Module,
      name_filter: Optional[Callable[[str], bool]] = None,
      checkpoint: Optional[Dict[str, TensorCheckpoint]] = None,
      **tkwargs: Any,
  ) -> Generator[Dict[str, TensorCheckpoint], None, None]:
      r"""Contextmanager that exits by rolling back a module's state_dict.
  
      Args:
          module: Module instance.
          name_filter: Optional Boolean function used to filter items by name.
          checkpoint: Optional cache of values and tensor metadata specifying the rollback
              state for the module (or some subset thereof).
          **tkwargs: Keyword arguments passed to `torch.Tensor.to` when copying data from
              each tensor in `module.state_dict()` to the internally created checkpoint.
              Only adhered to when the `checkpoint` argument is None.
  
      Yields:
          A dictionary of TensorCheckpoints for the module's state_dict. Any in-places
          changes to the checkpoint will be observed at rollback time. If the checkpoint
          is cleared, no rollback will occur.
      """
      # Create copies of the orginal values
      if checkpoint is None:
          checkpoint = {
              name: TensorCheckpoint(
                  values=values.detach().to(**tkwargs).clone(),
                  device=values.device,
                  dtype=values.dtype,
              )
              for name, values in module.state_dict().items()
              if name_filter is None or name_filter(name)
          }
  
      try:  # yield the checkpoint dictionary to the user
          yield checkpoint
      finally:  # restore original values of tracked parameters
          if checkpoint:
              state_dict = module.state_dict()
              for key, (values, device, dtype) in checkpoint.items():
                  tnsr = state_dict.get(key)
                  if tnsr is None:
                      state_dict[key] = values.to(
                          device=device, dtype=dtype
                      )
                  else:
                      tnsr[...] = values.to(device=device, dtype=dtype)
  
              module.load_state_dict(state_dict)
  
file_path: botorch/botorch/utils/context_managers.py
imports: from __future__ import annotations
  from contextlib import contextmanager
  from typing import Any, Callable, Dict, Generator, Iterable, NamedTuple, Optional, Union
  from torch import device, dtype, Tensor
  from torch.nn import Module
label_desc: Use image node_id 4 for calling a global function with example usage: module_rollback_ctx(module, name_filter, checkpoint) without return types