fix backward bug for custom device (#98586)
Fixes #ISSUE_NUMBER
In the backward on some device , it may get an error to get device index because of exchange a new thread.
So just set_device and check the device index in `setDevice` func may be better for some many kinds of devices.
For CUDA, the device index check is also included in `setDevice` func.https://github.com/pytorch/pytorch/blob/master/c10/cuda/impl/CUDAGuardImpl.h#:~:text=%7D-,void%20setDevice(Device%20d)%20const%20override%20%7B,%7D,-void%20uncheckedSetDevice(Device
```
void setDevice(Device d) const override {
TORCH_INTERNAL_ASSERT(d.is_cuda());
Device current_device = getDevice();
if (current_device != d) {
C10_CUDA_CHECK(cudaSetDevice(d.index()));
}
}
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/98586
Approved by: https://github.com/albanD
