Fix KeyError: 'aten::reshape' on calling torch.reshape (#3341)

microsoft · Feb 2, 2021 · be3a696 · be3a696
1 parent 533f2ef
commit be3a696
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/nni/retiarii/converter/op_types.py b/nni/retiarii/converter/op_types.py
@@ -29,6 +29,7 @@ class OpTypeName(str, Enum):
     'aten::cat': 'Cat',
     'aten::size': 'Size',
     'aten::view': 'View',
+    'aten::reshape': 'Reshape',
     'aten::eq': 'Eq',
     'aten::Bool': 'Bool',
     'aten::empty': 'Empty',

diff --git a/nni/retiarii/operation.py b/nni/retiarii/operation.py
@@ -150,6 +150,9 @@ def to_forward_code(self, field: str, output: str, inputs: List[str]) -> str:
         elif self.type == 'aten::view':
             assert len(inputs) == 2
             return f'{output} = {inputs[0]}.view({inputs[1]})'
+        elif self.type == 'aten::reshape':
+            assert len(inputs) == 2
+            return f'{output} = {inputs[0]}.reshape({inputs[1]})'
         elif self.type == 'aten::slice':
             raise RuntimeError('not supposed to have aten::slice operation')
         elif self.type == 'aten::Bool':