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

improve annotation device parameters where a device ordinal is allowed (#113647)

Browse Source 
Using mypy in code that depends on pytorch, I noticed that the type annotation doesn't allow a device ordinal.

`error: Argument "device" to "to_empty" of "Module" has incompatible type "int"; expected "str | device"  [arg-type]`

Pull Request resolved: https://github.com/pytorch/pytorch/pull/113647
Approved by: https://github.com/albanD
This commit is contained in:
Brian Vaughan
2023-11-16 23:07:02 +00:00
committed by PyTorch MergeBot
parent a56af02913
commit dbb96ef30d
9 changed files with 53 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
torchgen/api/python.py
View File

@ -927,7 +927,7 @@ def argument_type_str_pyi(t: Type) -> str:
elif t.name == BaseTy.Layout:
ret = "_layout"
elif t.name == BaseTy.Device:
ret = "Union[_device, str, None]"
ret = "Optional[DeviceLikeType]"
elif t.name == BaseTy.MemoryFormat:
ret = "memory_format"
elif t.name == BaseTy.Dimname: