Fixes an error that no CUDA devices are available.
In fact, CUDA might be available on the scheduler machine, and because this part of the code was run by the scheduler, torch.cuda.is_available() returned True, which caused the job to crash. This problem is solved by applying the modification in this merge request.