Restore TorchBench PR tests (#2636)
Summary:
With the Docker image back https://github.com/pytorch/benchmark/pull/2635, we can properly run tests now. All failures are fixed with the exception of the memory leak issue when training `moco` which I'm not sure and will keep it as a TODO
Pull Request resolved: https://github.com/pytorch/benchmark/pull/2636
Reviewed By: malfet
Differential Revision: D80882141
Pulled By: huydhn
fbshipit-source-id: 9f055bd3daaf8c7aafb939da7b74ace4c6cfa8ea