Enable CUDA graph train on torchvision models (#706)
Summary:
Without CUDA Graph (with batch size 32):
```
$ python run.py resnet50 -t train -d cuda
Running train method from resnet50 on cuda in eager mode.
GPU Time: 1034.030 milliseconds
CPU Dispatch Time: 1026.865 milliseconds
CPU Total Wall Time: 1034.011 milliseconds
```
With CUDA Graph (with batch size 32):
```
$ python run.py resnet50 -t train -d cuda --train_cudagraph
Running train method from resnet50 on cuda in eager mode.
GPU Time: 1038.927 milliseconds
CPU Dispatch Time: 346.313 milliseconds
CPU Total Wall Time: 1038.941 milliseconds
```
# Latency by batch size (Train, fp32, on V100)
<google-sheets-html-origin>
Batch Size | Eager (ms) | CUDA Graph (ms) | Speedup
-- | -- | -- | --
1 | 89.033 | 50.233 | 43.58%
2 | 92.854 | 56.13 | 39.55%
4 | 93.676 | 71.465 | 23.71%
8 | 105.099 | 105.381 | -0.27%
16 | 167.292 | 167.966 | -0.40%
32 | 297.315 | 297.989 | -0.23%
64 | 561.262 | 562.029 | -0.14%
Pull Request resolved: https://github.com/pytorch/benchmark/pull/706
Reviewed By: ngimel
Differential Revision: D33720966
Pulled By: xuzhao9
fbshipit-source-id: 8d422d597a879488d14361466172d32a1eeb1f19