Fix mypy typing check for test_dataset (#50108)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/50108
Test Plan: Imported from OSS
Reviewed By: glaringlee
Differential Revision: D25789184
Pulled By: ejguan
fbshipit-source-id: 0eeeeeda62533e7137d56f313b7bf11406b32611
