fix PR #2428 (#2433)
Summary:
https://github.com/pytorch/benchmark/issues/2428 was code reviewed internally on fbsource, and was not tested on OSS. This PR fixes the changes during code review. Before importing to Phabricator, it has been code reviewed on GitHub and tested locally.
Pull Request resolved: https://github.com/pytorch/benchmark/pull/2433
Reviewed By: embg
Differential Revision: D61717464
Pulled By: sfzhu93
fbshipit-source-id: f75fa3081840baa9079c6fbef9c22baaf469bb7a
