def METHOD_NAME(self):
    class MyModel(torch.nn.Module):
        def __init__(self, raise_in_recomp):
            super().__init__()
            self.fwd_count = 0
            self.raise_in_recomp = raise_in_recomp
            self.a = torch.nn.Linear(2, 2)
        def forward(self, x):
            if self.raise_in_recomp and self.fwd_count == 1:
                raise RuntimeError("foo")
            else:
                if not self.raise_in_recomp:
                    # raise in the first forward
                    raise RuntimeError("foo")
                self.fwd_count += 1
                return self.a(x)
    m = MyModel(raise_in_recomp=True)
    m_seq = torch.nn.Sequential(OrderedDict({"m": m}))
    checkpoint(m_seq.m)
    inp = torch.randn(1, 2)
    out = m_seq(inp).sum()
    # Should raise in forward recomputation
    with self.assertRaisesRegex(RuntimeError, "foo"):
        out.backward()
    # Check that _ac_generator is cleared out
    self.assertEqual(None, checkpoint.state(m)._ac_generator)
    m = MyModel(raise_in_recomp=False)
    checkpoint(m)
    inp = torch.randn(1, 2)
    # Should raise in first forward
    with self.assertRaises(RuntimeError):
        m(inp)
    self.assertEqual(None, checkpoint.state(m)._ac_generator)