enforce `dtype` (reland) (#102996)
Summary: The original diff didn't break the test.
Test Plan: N/A
Differential Revision: D46448488
Pull Request resolved: https://github.com/pytorch/pytorch/pull/102996
Approved by: https://github.com/malfet, https://github.com/wanchaol
