Fix in benchmark_test_generator (#20237)
Summary:
Add missing import
Pull Request resolved: https://github.com/pytorch/pytorch/pull/20237
Differential Revision: D15245957
Pulled By: ilia-cher
fbshipit-source-id: 0f71aa08eb9ecac32002a1644838d06ab9faa37c