[JIT] Freeze allows preservation of submodule attributes (#66102)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/66102
This changes allows the `preserved_attributes` parameter of `torch.jit.freeze` to accept attributes of submodules. Previously, only root-level attributes were able to be preserved. Example:
```
class SubModule(nn.Module):
def __init__(self):
super(SubModule, self).__init__()
self.a = 1
self.b = 2
def forward(self):
return self.a + self.b
class Module(nn.Module):
def __init__(self):
super(Module, self).__init__()
self.sub = SubModule()
def forward(self):
return self.sub()
mod = torch.jit.script(Module())
mod.eval()
frozen_mod = torch.jit.freeze(mod, preserved_attrs = ['sub.a'])
mod.sub # OK
mod.sub.a # OK
mod.sub.b # Error, not preserved
mod() # = 3
mod.sub.a = 0
mod() # = 2
```
Test Plan: `buck test caffe2/test:jit -- TestFreezing`
Reviewed By: eellison
Differential Revision: D31383868
fbshipit-source-id: 34a05ca9528d4e5f04f71ac2a339fd584a8fa305
pytorch
dd81fa90 - [JIT] Freeze allows preservation of submodule attributes (#66102)
