Select GPU device through env vars

1

I’ve seen that the GPU device can be picked by passing an argument to the cuda method.
However, I was wondering if it would be possible to define this through environment variables.

4 Likes
2

Yes, it is possible.
You can either run your script by setting CUDA_VISIBLE_DEVICES like

CUDA_VISIBLE_DEVICES=1 python myscript.py

Or change it in the code, as done in this script

12 Likes
3

thanks! are there any docs for this? I really couldn’t find it :-s

4

This is specific to CUDA, and not to pytorch.
You can find some more information about it in this link

2 Likes
5

Ah, I understand why I couldn’t find it in the docs. :slight_smile: thanks again.

6

One caveat for setting the CUDA_VISIBLE_DEVICES environment variable in the code is that that has to be done before any call to torch.cuda.

4 Likes