Skip to content

Commit f512d7e

Browse files
authored
Arm backend: Clear const shapes cache before/after use (pytorch#19914)
Signed-off-by: Martin Lindström <Martin.Lindstroem@arm.com>
1 parent feb84f8 commit f512d7e

1 file changed

Lines changed: 7 additions & 0 deletions

File tree

backends/arm/_passes/insert_const_shapes.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,13 @@ def _is_shape_arg(arg: Any) -> bool:
4040
and all(isinstance(x, int) for x in arg)
4141
)
4242

43+
def call(self, graph_module):
44+
self._const_shape_cache.clear()
45+
try:
46+
return super().call(graph_module)
47+
finally:
48+
self._const_shape_cache.clear()
49+
4350
def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False):
4451
if op not in self.target_ops:
4552
return super().call_operator(op, args, kwargs, meta, updated)

0 commit comments

Comments
 (0)