It is pretty strange but the $CUDA_VISIBLE_DEVICES is not in my environment when I execute the code, the code is executed via slurm and nvidia module is loaded correctly as well as the GPU but the variable is not setted, is there a way to set it up by myself?