saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 17:15:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Override ORTModule named_modules to support extra arg (#7954)

Browse source
This commit is contained in:
baijumeswani 2021-06-08 17:32:09 -07:00 committed by GitHub
parent 1cc896c8ae
commit f38200e209
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
orttraining/orttraining/python/training/ortmodule/ortmodule.py
View file

@ -182,10 +182,12 @@ class ORTModule(torch.nn.Module):
yield from self._module_metadata.original_module.modules()
def named_modules(self, memo: Optional[Set['Module']] = None, prefix: str = ''):
def named_modules(self, *args, **kwargs):
"""Override original method to delegate execution to the original PyTorch user module"""
yield from self._module_metadata.original_module.named_modules(memo, prefix)
# PyTorch >1.8.1 has an extra arg remove_duplicate that is not present in 1.8.1
# To support both, use args and kwargs (since user can call the method with only positional args or kwargs)
yield from self._module_metadata.original_module.named_modules(*args, **kwargs)
def add_module(self, name: str, module: Optional['Module']) -> None:
"""Raises a NotImplementedError exception since ORTModule does not support adding modules to it"""