Skip to content

Commit

Permalink
Showing 1 changed file with 15 additions and 0 deletions.
15 changes: 15 additions & 0 deletions python/tvm/relay/frontend/pytorch.py
Original file line number Diff line number Diff line change
Expand Up @@ -3402,6 +3402,21 @@ def _handel_nested_input(inputs):
relay_out = relay_op(
inputs, _get_input_types(op_node, outputs, default_dtype=self.default_dtype)
)

def do_mutate(sym):
class PureMutator(ExprMutator):
def __init__(self):
ExprMutator.__init__(self)

def mutate(self, sym):
if isinstance(sym, _expr.TupleWrapper):
return _expr.TupleWrapper(self.visit(sym.tuple_value), sym.size)
if isinstance(sym, _expr.RelayExpr):
return self.visit(sym)
return sym
return PureMutator().mutate(sym)

relay_out = do_mutate(relay_out)
self.record_output_type(relay_out)

if isinstance(relay_out, tuple):
Expand Down

0 comments on commit e9b65c7

Please sign in to comment.