Open
Description
🐛 Describe the bug
When using typing.get_type_hints on a TorchScript model’s forward method, it works before saving and loading, but fails after loading the scripted model.
import sys
from typing import get_type_hints
import torch
import torch.nn as nn
print("Python version:", sys.version)
print("PyTorch version:", torch.__version__)
class DummyModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.linear(x)
model1 = DummyModel()
model1 = torch.jit.script(model1)
print("########## Model 1 ##########")
print(get_type_hints(model1.forward))
model1.save("dummy_model.pt")
model2 = torch.jit.load("dummy_model.pt")
model2 = torch.jit.script(model2)
print("########## Model 2 ##########")
print(get_type_hints(model2.forward))Output:
Python version: 3.10.16 (main, Feb 12 2025, 14:50:02) [Clang 19.1.6 ]
PyTorch version: 2.7.0+cu126
########## Model 1 ##########
{'x': <class 'torch.Tensor'>, 'return': <class 'torch.Tensor'>}
########## Model 2 ##########
Traceback (most recent call last):
File ".../deploy.py", line 32, in <module>
print(get_type_hints(model2.forward))
File ".../.local/share/uv/python/cpython-3.10.16-linux-x86_64-gnu/lib/python3.10/typing.py", line 1856, in get_type_hints
raise TypeError('{!r} is not a module, class, method, '
TypeError: <torch.ScriptMethod object at 0x7f97dfde3a10> is not a module, class, method, or function.
Versions
Python version: 3.10.16 (main, Feb 12 2025, 14:50:02) [Clang 19.1.6 ]
PyTorch version: 2.7.0+cu126
I'm running the script with uv, versions have been printed above.
Command to run:
uv run --no-project --with torch==2.7.0 script.py