| 145 | inputs = eager_module.get_random_inputs() # type: ignore[operator] |
| 146 | |
| 147 | class WrapperModule(torch.nn.Module): |
| 148 | def __init__(self, fn, method_name=method_name): |
| 149 | super().__init__() |
| 150 | self.fn = fn |
| 151 | self.method_name = method_name |
| 152 | |
| 153 | def forward(self, *args, **kwargs): |
| 154 | return getattr(self.fn, self.method_name)(*args, **kwargs) |
| 155 | |
| 156 | if method_name != "forward": |
| 157 | # Only require wrapper module if we're exporting a specific method other than forward. |
no outgoing calls
no test coverage detected