mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
83d88d128d
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: D72986826 Pull Request resolved: https://github.com/pytorch/pytorch/pull/151264 Approved by: https://github.com/angelayi
11 lines
271 B
Python
11 lines
271 B
Python
import torch
|
|
|
|
|
|
class _WrapperModule(torch.nn.Module):
|
|
def __init__(self, f): # type: ignore[no-untyped-def]
|
|
super().__init__()
|
|
self.f = f
|
|
|
|
def forward(self, *args, **kwargs): # type: ignore[no-untyped-def]
|
|
return self.f(*args, **kwargs)
|