Skip to content

Commit

Permalink
Skip Removal of squeeze which contain dynamic shapes
Browse files Browse the repository at this point in the history
  • Loading branch information
meenakshiramanathan1 committed Jan 17, 2025
1 parent 38d2522 commit 8f67e84
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions python/tvm/relay/op/contrib/forge/forge_passes.py
Original file line number Diff line number Diff line change
Expand Up @@ -1883,6 +1883,9 @@ def __init__(self):
def callback(self, pre, post, node_map):
act = node_map[self.act][0]
axis = post.attrs.axis
# Skip removal of squeeze which contain dynamic shapes
if any([isinstance(dim, tvm.tir.expr.Any) for dim in pre.checked_type.shape]):
return post
input_shape = [int(dim) for dim in pre.args[0].checked_type.shape]
adjusted_axes = [(ax - len(input_shape)) if ax >= 0 else ax for ax in axis]
assert all(ax < 0 for ax in adjusted_axes), "Invalid squeeze dimension: all axes must be negative."
Expand Down

0 comments on commit 8f67e84

Please sign in to comment.