code: def test_module_rollback_ctx(self):
      # Test that only unfiltered objects get rolled back
      a = self.module.a.detach().clone()
      b = self.module.b.detach().clone()
      c = self.module.c.detach().clone()
      with module_rollback_ctx(
          self.module, lambda name: name == "a", dtype=torch.float16
      ) as ckpt:
          for (
              tnsr,
              _,
              __,
          ) in ckpt.values():  # test whether dtype is obeyed
              self.assertEqual(torch.float16, tnsr.dtype)
  
          self.module.a.data[...] = 0
          self.module.b.data[...] = 0
          self.module.c.data[...] = 0
  
      self.assertTrue(self.module.a.equal(a))
      self.assertTrue(self.module.b.eq(0).all())
      self.assertTrue(self.module.c.eq(0).all())
  
      # Test that changes to checkpoint dict are reflected in rollback state
      with module_rollback_ctx(self.module) as ckpt:
          self.module.a.data[...] = 1
          self.module.b.data[...] = 1
          self.module.c.data[...] = 1
          del ckpt["a"]
  
      self.assertTrue(self.module.a.eq(1).all())
      self.assertTrue(self.module.b.eq(0).all())
      self.assertTrue(self.module.c.eq(0).all())
  
      # Test rolling back to a user-provided checkpoint
      checkpoint = {
          "a": TensorCheckpoint(a, a.device, a.dtype),
          "b": TensorCheckpoint(b, b.device, b.dtype),
          "c": TensorCheckpoint(c, c.device, c.dtype),
      }
      with module_rollback_ctx(
          module=self.module, checkpoint=checkpoint
      ):
          pass
      self.assertTrue(self.module.a.equal(a))
      self.assertTrue(self.module.b.equal(b))
      self.assertTrue(self.module.c.equal(c))
  
      # Test that items in checkpoint get inserted into state_dict
      with delattr_ctx(self.module, "a"):
          with self.assertRaisesRegex(  # should fail when attempting to rollback
              RuntimeError, r'Unexpected key\(s\) in state_dict: "a"'
          ):
              with module_rollback_ctx(
                  module=self.module, checkpoint=checkpoint
              ):
                  pass
  
file_path: botorch/test/utils/test_context_managers.py
imports: from __future__ import annotations
  from string import ascii_lowercase
  import torch
  from botorch.utils.context_managers import delattr_ctx, module_rollback_ctx, parameter_rollback_ctx, requires_grad_ctx, TensorCheckpoint, zero_grad_ctx
  from botorch.utils.testing import BotorchTestCase
  from torch.nn import Module, Parameter
label_desc: Use image node_id 5 for calling the TestContextManagers obj's underlying member method code with example usage: obj.test_module_rollback_ctx() without return types