mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
dad5e5e262
Summary: as title `export._trace._WrapperModule` is used to wrap functions into a Module so we can export the function. We add `export._wrapper_utils` to `dynamo`'s `MOD_INLINELIST` so dynamo traces into `_WrapperModule` Fixes https://github.com/pytorch/pytorch/issues/146867 Test Plan: ``` buck run fbcode//mode/dev-nosan //caffe2/test:test_export -- -r wrapper_module ``` Differential Revision: D69434316 Pull Request resolved: https://github.com/pytorch/pytorch/pull/146919 Approved by: https://github.com/angelayi
11 lines
207 B
Python
11 lines
207 B
Python
import torch
|
|
|
|
|
|
class _WrapperModule(torch.nn.Module):
|
|
def __init__(self, f):
|
|
super().__init__()
|
|
self.f = f
|
|
|
|
def forward(self, *args, **kwargs):
|
|
return self.f(*args, **kwargs)
|