Update module wrapper so that params are explicitly registered to the wrapper (#10357)

Pull Request resolved: #10305

Seeing issue with linear where the fqns for constants disappear.

Registering self.method_name as a submodule of wrapper means that the
parameters are registered to the wrapper. thanks @angelayi for the fix!

```
  File "/data/users/lfq/fbsource/buck-out/v2/gen/fbcode/1af94fa701700343/executorch/test/models/__export_delegated_program__/export_delegated_program#link-tree/torch/export/_trace.py", line 1980, in _export_for_training
    export_artifact = export_func(
  File "/data/users/lfq/fbsource/buck-out/v2/gen/fbcode/1af94fa701700343/executorch/test/models/__export_delegated_program__/export_delegated_program#link-tree/torch/export/_trace.py", line 1473, in _strict_export
    _replace_param_buffer_names(param_buffer_table, export_graph_signature)
  File "/data/users/lfq/fbsource/buck-out/v2/gen/fbcode/1af94fa701700343/executorch/test/models/__export_delegated_program__/export_delegated_program#link-tree/torch/export/_trace.py", line 272, in _replace_param_buffer_names
    spec.target = param_buffer_table[spec.target]
KeyError: 'L__self___fn___self___linear.weight'
```
ghstack-source-id: 279346028
@exported-using-ghexport

Differential Revision:
[D73279618](https://our.internmc.facebook.com/intern/diff/D73279618/)

Author: lucylq

test/models/export_delegated_program.py

@@ -90,6 +90,18 @@ def get_random_inputs(self) -> Sequence[torch.Tensor]:
         return (torch.ones(n, n, n), 2 * torch.ones(n, n, n), 3 * torch.ones(n, n, n))
 
 
+class ModuleLinear(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(3, 3)
+
+    def forward(self, x: torch.Tensor):
+        return self.linear(x)
+
+    def get_random_inputs(self):
+        return (torch.randn(3),)
+
+
 #
 # Backends
 #
@@ -116,24 +128,23 @@ def export_module_to_program(
     extract_delegate_segments: bool,
     constant_tensor_alignment: Optional[int] = None,
     delegate_alignment: Optional[int] = None,
-    method: str = "forward",
+    method_name: str = "forward",
 ) -> ExecutorchProgramManager:
     eager_module = module_class().eval()
     inputs = ()
     if hasattr(eager_module, "get_random_inputs"):
         inputs = eager_module.get_random_inputs()  # type: ignore[operator]
 
     class WrapperModule(torch.nn.Module):
-        def __init__(self, fn):
+        def __init__(self, fn, method_name=method_name):
             super().__init__()
             self.fn = fn
+            self.method_name = method_name
 
         def forward(self, *args, **kwargs):
-            return self.fn(*args, **kwargs)
+            return getattr(self.fn, self.method_name)(*args, **kwargs)
 
-    exported_program = export(
-        WrapperModule(getattr(eager_module, method)), args=inputs, strict=True
-    )
+    exported_program = export(WrapperModule(eager_module), args=inputs, strict=True)
 
     edge_config = EdgeCompileConfig(_check_ir_validity=False)
     et_config = exir.ExecutorchBackendConfig(

test/models/targets.bzl

@@ -156,6 +156,7 @@ def define_common_targets():
         "ModuleAddMul",
         "ModuleAddLarge",
         "ModuleSubLarge",
+        "ModuleLinear",
     ]
 
     # Name of the backend to use when exporting delegated programs.