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

[Dynamo] Support zip_longest (#131497)

Browse Source 
Fixes #121348

Pull Request resolved: https://github.com/pytorch/pytorch/pull/131497
Approved by: https://github.com/mlazos, https://github.com/jansel, https://github.com/zou3519
This commit is contained in:
Yanbo Liang
2024-07-26 14:06:10 +00:00
committed by PyTorch MergeBot
parent c9888c2739
commit e76e566cfb
3 changed files with 35 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
torch/_dynamo/variables/iter.py
View File

@ -196,6 +196,10 @@ class ItertoolsVariable(VariableTracker):
return variables.UserFunctionVariable(polyfill.dropwhile).call_function(
tx, args, kwargs
)
elif self.value is itertools.zip_longest:
return variables.UserFunctionVariable(polyfill.zip_longest).call_function(
tx, args, kwargs
)
else:
return super().call_function(tx, args, kwargs)