Fix pytorch_unet model (#1146)
Summary:
Fix the broken jit test for pytorch_unet
Pull Request resolved: https://github.com/pytorch/benchmark/pull/1146
Reviewed By: erichan1
Differential Revision: D39180853
Pulled By: xuzhao9
fbshipit-source-id: 9f1a3bc12d975d28a841c6c287063bfc3fd4e0d4
