Support auto generation of device check (#56872)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/56872
ghstack-source-id: 127914018
Test Plan: auto test
Reviewed By: ezyang
Differential Revision: D27986429
fbshipit-source-id: 0da8413b0b8e6810fcea27ed1de499f11f68bd1f
