Fix a major bug (#1606)
Summary:
I should have tested better
Pull Request resolved: https://github.com/pytorch/benchmark/pull/1606
Reviewed By: xuzhao9
Differential Revision: D45458183
Pulled By: janeyx99
fbshipit-source-id: 5c1ffd5484a92a0b40604433c65b54915f71c6e0
