Add ability to freeze storages inside functionalization (#88141)

Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/88141
Approved by: https://github.com/albanD, https://github.com/bdhirsh
This commit is contained in:
Edward Z. Yang
2022-10-31 14:42:19 -07:00
committed by PyTorch MergeBot
parent 61f955dd83
commit bb7e6254e4
6 changed files with 53 additions and 0 deletions

View File

@ -102,6 +102,18 @@ class TestFunctionalization(TestCase):
return z2
self.assert_functionalization(f, torch.ones(4))
def test_freeze(self):
def f(x):
y = x.clone()
z = y[0]
torch._freeze_functional_tensor(y)
x.add_(1)
self.assertRaises(RuntimeError, lambda: y.add_(1))
self.assertRaises(RuntimeError, lambda: z.add_(1))
return z
_functionalize(f, reapply_views=True)(torch.ones(3, 3))
def test_view_clone_view_inplace(self):
def f(input):
shape = [1, 1024, 128, 128]