Fix handling of device (#78615)
Removes an unnecessary auxiliary function (we had already implemented
it), uses DeviceLikeType to denote str or dtype, and adds `is_cpu` and
`is_cuda` helper functions
Pull Request resolved: https://github.com/pytorch/pytorch/pull/78615
Approved by: https://github.com/mruberry