Changing return of relay.Function or IRModule?

Try this:

        new_body = self.visit(fn.body)
        print("Visited all", new_body)
        return_values_to_function = relay.Tuple([new_body] + self.return_values)
        func = relay.Function(fn.params, return_values_to_function, fn.ret_type, fn.type_params, fn.attrs)