Rewrite all inputs to Relay IR (convert 1d to 4D)

Hi @mbrookhart, thanks.

I got it working by replacing the last line with func = relay.Function(analysis.free_vars(new_body), new_body)

Not sure, why the one that I originally used not working. But this solves the problem. Thanks.