[torchscript] Add logging for model id. (#130118)

Summary: as title.

Test Plan: CI

Reviewed By: angelayi

Differential Revision: D59348256

Pull Request resolved: https://github.com/pytorch/pytorch/pull/130118
Approved by: https://github.com/BoyuanFeng
This commit is contained in:
Zhengxu Chen
2024-07-09 22:24:16 +00:00
committed by PyTorch MergeBot
parent fb5cb17fbe
commit 37d4d04309
5 changed files with 27 additions and 14 deletions

View File

@ -1520,6 +1520,15 @@ def _extract_tensors(obj):
return tensors
def _get_model_id(obj) -> Optional[str]:
if isinstance(obj, torch.jit.ScriptModule):
return str(obj._c._type())
elif isinstance(obj, torch.jit.ScriptFunction):
return obj.qualified_name
else:
return None
# In Python-3.11+ typed enums (i.e. IntEnum for example) retain number of base class methods in subclass
# that were previously dropped. To preserve the behavior, explicitly drop them there