Repository: pytorch
Issue: [JIT] magic methods do not work after reloading model
Labels: oncall: jit

### 🐛 Describe the bug

```python
import torch
from torch import Tensor, jit, nn

class Foo(nn.Module):
    weight: Tensor

    def __init__(self, m: int, n: int):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(m, n))

    @jit.export
    def __len__(self) -> int:
        return len(self.weight)

    @jit.export
    def length(self) -> int:
        return len(self.weight)

model = Foo(2, 3)

print(model.length())  # works! ✓
print(len(model))  # works! ✓

model = jit.script(model)  # scripting the model

print(model.length())  # works! ✓
print(len(model))  # works! ✓

jit.save(model, "model.pt")  # saving and reloading the model
model = jit.load("model.pt")

print(model.length())  # works! ✓
print(len(model))  # doesn't work 😢
```

It would be useful if this worked, since certain models have e.g. (re-) initialization schemes that depend on the depth of the model.

### Versions

<details><summary>Environment</summary>

```
Collecting environment information...
PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 21.10 (x86_64)
GCC version: (Ubuntu 11.2.0-7ubuntu2) 11.2.0
Clang version: 13.0.0-2
CMake version: Could not collect
Libc version: glibc-2.34

Python version: 3.9.12 | packaged by conda-forge | (main, Mar 24 2022, 23:25:59)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.13.0-40-generic-x86_64-with-glibc2.34
Is CUDA available: False
CUDA runtime version: 11.5.119
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.4.0
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.1
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn.so.8.2.4
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.4
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.4
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.4
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.4
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.4
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.4
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.942
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.6
[pip3] numpydoc==1.3.1
[pip3] pytorch-forecasting==0.10.1
[pip3] pytorch-lightning==1.6.2
[pip3] torch==1.11.0
[pip3] torch-tb-profiler==0.4.0
[pip3] torchaudio==0.11.0
[pip3] torchdata==0.4.0a0+b739701
[pip3] torchdiffeq==0.2.2
[pip3] torchinfo==1.6.5
[pip3] torchmetrics==0.8.1
[pip3] torchtext==0.12.0
[pip3] torchvision==0.12.0
[conda] blas                      2.114                       mkl    conda-forge
[conda] blas-devel                3.9.0            14_linux64_mkl    conda-forge
[conda] cudatoolkit               11.3.1               ha36c431_9    nvidia
[conda] libblas                   3.9.0            14_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            14_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            14_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            14_linux64_mkl    conda-forge
[conda] mkl                       2022.0.1           h06a4308_117  
[conda] mkl-devel                 2022.0.1           h66538d2_117  
[conda] mkl-include               2022.0.1           h06a4308_117  
[conda] mypy                      0.942            py39hb9d737c_0    conda-forge
[conda] mypy_extensions           0.4.3            py39h06a4308_1  
[conda] numexpr                   2.8.0           mkl_py39h25e7801_2    conda-forge
[conda] numpy                     1.21.6           py39h18676bf_0    conda-forge
[conda] numpydoc                  1.3.1              pyhd8ed1ab_0    conda-forge
[conda] pytorch                   1.11.0          py3.9_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-forecasting       0.10.1             pyhd8ed1ab_0    conda-forge
[conda] pytorch-lightning         1.6.2              pyhd8ed1ab_0    conda-forge
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch-tb-profiler         0.4.0                    pypi_0    pypi
[conda] torchaudio                0.11.0               py39_cu113    pytorch
[conda] torchdiffeq               0.2.2              pyhd8ed1ab_0    conda-forge
[conda] torchinfo                 1.6.5              pyhd8ed1ab_0    conda-forge
[conda] torchmetrics              0.8.1              pyhd8ed1ab_0    conda-forge
[conda] torchtext                 0.12.0                     py39    pytorch
[conda] torchvision               0.12.0               py39_cu113    pytorch
```

</details>