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
Reviewed By: malfet
Differential Revision: D80882141
Pulled By: huydhn
