[DataLoader] Fix bugs for typing (#58450)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/58450
Test Plan: Imported from OSS
Reviewed By: VitalyFedyunin
Differential Revision: D28507877
Pulled By: ejguan
fbshipit-source-id: f4051ff51ce77ef45214f11cba10c8a7e1da4dad
