Repository: pytorch
Issue: torch.jit.script: scripting doesn't work with wraps
Labels: oncall: jit

### 🐛 Describe the bug

Say we have a decorator that does some kind of preprocessing (maybe some checks on the arguments, maybe some kind of normalization etc) and we decide to decorate and then script a function as shown below:

```python
import functools
import torch


def decorator(obj):
    @functools.wraps(obj)
    def wrapper(inpt):
        return obj((inpt).sigmoid())
    return wrapper


@decorator
def func(inpt):
    return torch.sqrt(inpt)


script_func = torch.jit.script(func)

inpt = torch.ones(1)
print(func(inpt))
print(script_func(inpt))
print(script_func.code)

```
Running this gives:

```
tensor([0.8550])
tensor([1.])
def func(inpt: Tensor) -> Tensor:
  return torch.sqrt(inpt)
```

So it seems that scripting jumps straight to `__wrapped__` and ignores the decorator.

If we separate the decorator and the function:

```python
# decorator_file.py
import functools


def decorator(obj):
    @functools.wraps(obj)
    def wrapper(inpt):
        return obj((inpt).sigmoid())
    return wrapper

```

```python
# script_file.py
import torch
from decorator_file import decorator


@decorator
def func(inpt):
    return torch.sqrt(inpt)


script_func = torch.jit.script(func)

```

We get the following:

```
Traceback (most recent call last):
  File "C:\Users\user\Desktop\test-env\script_file.py", line 10, in <module>
    script_func = torch.jit.script(func)
                  ^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\user\Desktop\test-env\env\Lib\site-packages\torch\jit\_script.py", line 1341, in script
    fn = torch._C._jit_script_compile(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError:
undefined value torch:
  File "C:\Users\user\Desktop\test-env\decorator_file.py", line 7
@decorator
def func(inpt):
    return torch.sqrt(inpt)
           ~~~~~ <--- HERE
```

So the decorator and the imports from `script_file.py` are ignored (need to `import torch` in `decorator_file.py`).

### Versions

```
Collecting environment information...
PyTorch version: 2.0.1+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Home
GCC version: (x86_64-posix-seh-rev0, Built by MinGW-W64 project) 8.1.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.3 (tags/v3.11.3:f3909b8, Apr  4 2023, 23:49:59) [MSC v.1934 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19045-SP0
Is CUDA available: False
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1650
Nvidia driver version: 471.11
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=2592
DeviceID=CPU0
Family=198
L2CacheSize=1536
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=2592
Name=Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz
ProcessorType=3
Revision=

Versions of relevant libraries:
[pip3] torch==2.0.1
[conda] Could not collect

```

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel