Simplify device guard code generation (#55112)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/55112
Based on https://github.com/pytorch/pytorch/pull/47765
ghstack-source-id: 126114775
Test Plan: buck build //caffe2/aten/...
Reviewed By: ezyang
Differential Revision: D27487085
fbshipit-source-id: 157fcd19f538ce0c1e053e3e974b48bdb93a0226
