Fix code review from #48659 and #48116 (#48731)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48731
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Test Plan: Imported from OSS
Reviewed By: bhosmer
Differential Revision: D25278034
Pulled By: ezyang
fbshipit-source-id: 73652311b48d8d80c06e9385b7ff18ef3a158ae8
