Fix dispatching issue of the new device type. (#97273)
Summary: Fix the device type dispatching issue.
Test Plan: All CI should pass.
Reviewed By: scottxu0730
Differential Revision: D44179512
Pull Request resolved: https://github.com/pytorch/pytorch/pull/97273
Approved by: https://github.com/ezyang
