frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[BE] Improve the typing related to model input argument of torch.compile() (#153559)

Browse Source 
Summary: Match the `overload` typing with the original typing in function definition and adjust the corresponding comments.

Test Plan: contbuild & OSS CI

Differential Revision: D74746243

Pull Request resolved: https://github.com/pytorch/pytorch/pull/153559
Approved by: https://github.com/Skylion007
This commit is contained in:
Tsung-Hsien Lee
2025-05-15 04:49:26 +00:00
committed by PyTorch MergeBot
parent d2f6c6df1d
commit 756fd80734
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
torch/__init__.py
View File

@ -2448,7 +2448,7 @@ def compile(
def compile(
model: _Optional[_Callable] = None,
model: _Optional[_Callable[_InputT, _RetT]] = None,
*,
fullgraph: builtins.bool = False,
dynamic: _Optional[builtins.bool] = None,
@ -2479,7 +2479,7 @@ def compile(
function, they will all share the same code cache.
Args:
model (Callable): Module/function to optimize
model (Callable or None): Module/function to optimize
fullgraph (bool): If False (default), torch.compile attempts to discover compileable regions
in the function that it will optimize. If True, then we require that the entire function be
capturable into a single graph. If this is not possible (that is, if there are graph breaks),