mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
dccc1ca839
Summary: This fixes a case left incomplete by https://github.com/pytorch/pytorch/pull/106229 The object is using __prepare_scriptable__ correctly inside of torch.jit.script() but the clousre that is obtained below is using the non-prepared version. This causes issues when the prepared and non-prepared versions are in different python modules. Test Plan: ``` buck2 run mode/opt caffe2/test:jit -- -r test_decorator ``` Differential Revision: D54308741 Re-exporting, as #120806 #121307 were not properly merged. Co-authored-by: Daniel Herrera <dherrera@meta.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/121553 Approved by: https://github.com/huydhn, https://github.com/seemethere
21 lines
529 B
Python
21 lines
529 B
Python
r"""
|
|
Decorator used in test_decorator.py. We define it in a
|
|
separate file on purpose to test that the names in different modules
|
|
are resolved correctly.
|
|
"""
|
|
|
|
import functools
|
|
|
|
|
|
def my_decorator(func):
|
|
"""Dummy decorator that removes itself when torchscripting"""
|
|
|
|
@functools.wraps(func)
|
|
def wrapped_func(*args, **kwargs):
|
|
return func(*args, **kwargs)
|
|
|
|
# torch.jit.script() uses __prepare_scriptable__ to remove the decorator
|
|
wrapped_func.__prepare_scriptable__ = lambda: func
|
|
|
|
return wrapped_func
|