text: import torch
  from typing import Any
  
  def set_saved_tensors_default_hooks(pack_hook, unpack_hook):
      torch._C._autograd._register_default_hooks(pack_hook, unpack_hook)
  
  def reset_saved_tensors_default_hooks():
      torch._C._autograd._reset_default_hooks()
  
  class save_on_cpu(object):
      def __init__(self, pin_memory=False):
          def pack_to_cpu(tensor):
              if not pin_memory:
                  return (tensor.device, tensor.cpu())
  
              storage = torch.empty(
                  tensor.size(),
                  dtype=tensor.dtype,
                  layout=tensor.layout,
                  pin_memory=(torch.cuda.is_available() and not tensor.is_sparse))
              storage.copy_(tensor)
              return (tensor.device, storage)
  
          def unpack_from_cpu(packed):
              device, tensor = packed
              return tensor.to(device, non_blocking=pin_memory)
  
          self.pack_hook = pack_to_cpu
          self.unpack_hook = unpack_from_cpu
  
      def __enter__(self):
          torch._C._autograd._register_default_hooks(self.pack_hook, self.unpack_hook)
  
      def __exit__(self, *args: Any):
          torch._C._autograd._reset_default_hooks()
  