Repository: pytorch
Issue: RuntimeError:  Expression of type - cannot be used in a type expression: __torch__.transformers_modules.code-5p-110m-embedding.modeling_codet5p_embedding.___torch_mangle_1368.CodeT5pEmbeddingModel ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
Labels: oncall: jit

### 🐛 Describe the bug

I use the model: https://huggingface.co/Salesforce/codet5p-110m-embedding and with `torch.jit.load`, it makes error:
```python
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel

model = AutoModel.from_pretrained('Salesforce/codet5p-110m-embedding', trust_remote_code=True)
input_list = [tokenizer(["xxxxxx"], padding='max_length', truncation=True, max_length=100, return_tensors='pt')['input_ids']]
ts_model = torch.jit.trace(model.eval(), input_list)
torch.jit.save(ts_model, "codet5p_embedding.pt")
ts_model = torch.jit.load("codet5p_embedding.pt").cuda()
```

the error message is:
```
RuntimeError                              Traceback (most recent call last)
Cell In[21], line 1
----> 1 ts_model = torch.jit.load("codet5p_embedding.pt").cuda()
      2 org_out = model(input_list)
      3 ts_out = ts_model(input_list)

File "/opt/conda/lib/python3.10/site-packages/torch/jit/_serialization.py", line 162, in load(f, map_location, _extra_files, _restore_shapes)
    160 cu = torch._C.CompilationUnit()
    161 if isinstance(f, (str, pathlib.Path)):
--> 162     cpp_module = torch._C.import_ir_module(cu, str(f), map_location, _extra_files, _restore_shapes)  # type: ignore[call-arg]
    163 else:
    164     cpp_module = torch._C.import_ir_module_from_buffer(
    165         cu, f.read(), map_location, _extra_files, _restore_shapes
    166     )  # type: ignore[call-arg]

RuntimeError: 
Expression of type - cannot be used in a type expression:
__torch__.transformers_modules.code-5p-110m-embedding.modeling_codet5p_embedding.___torch_mangle_1368.CodeT5pEmbeddingModel
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
```

### Versions

Collecting environment information...
PyTorch version: 2.1.2
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 4.8.5 20150623 (Red Hat 4.8.5-44)
Clang version: Could not collect
CMake version: version 3.28.4
Libc version: glibc-2.17

Python version: 3.10.10 (main, Mar 21 2023, 18:45:11) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.10.112-005.ali5000.al8.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA L20
Nvidia driver version: 535.161.08
cuDNN version: Probably one of the following:
/usr/lib64/libcudnn.so.8.9.0
/usr/lib64/libcudnn_adv_infer.so.8.9.0
/usr/lib64/libcudnn_adv_train.so.8.9.0
/usr/lib64/libcudnn_cnn_infer.so.8.9.0
/usr/lib64/libcudnn_cnn_train.so.8.9.0
/usr/lib64/libcudnn_ops_infer.so.8.9.0
/usr/lib64/libcudnn_ops_train.so.8.9.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.1.2
[pip3] torch-dct==0.1.6
[pip3] torchaudio==2.1.2
[pip3] torchvision==0.16.2
[pip3] triton==2.0.0
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] mkl-service               2.4.0           py310h5eee18b_1  
[conda] mkl_fft                   1.3.8           py310h5eee18b_0  
[conda] mkl_random                1.2.4           py310hdb19cb5_0  
[conda] numpy                     1.26.4          py310h5f9d8c6_0  
[conda] numpy-base                1.26.4          py310hb5e798b_0  
[conda] pytorch                   2.1.2           py3.10_cuda12.1_cudnn8.9.2_0    pytorch
[conda] pytorch-cuda              12.1                 ha16c6d3_5    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch-dct                 0.1.6                    pypi_0    pypi
[conda] torchaudio                2.1.2               py310_cu121    pytorch
[conda] torchvision               0.16.2              py310_cu121    pytorch
[conda] triton                    2.0.0                    pypi_0    pypi

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel