Specifying GPU in IPython notebook

I have found this piece of code :

import os
os.environ[“CUDA_DEVICE_ORDER”]=“PCI_BUS_ID”
os.environ[“CUDA_VISIBLE_DEVICES”]=“1”

from this website . Did anybody tried to use gpu’s programmatically inside a notebook. Does the code looks valid?

I have tried to use these environment variables inside a notebook and command line.

We use this when our CUDA application need to target a specific GPU. Note, be careful, they don’t mean PyTorch application here, but still works.

To use it, set CUDA_VISIBLE_DEVICES to a comma-separated list of device IDs to make only those devices visible to the application. So, your code is valid.

These CUDA APIs are much more low level way of controlling the GPU(s). For the majority of us building PyTorch app, the ‘recommended’ way to achieve something similar is to use PyTorch API: torch.cuda.set_device(0).

1 Like

Is there anything similar to torch.cuda.set_device(0) in tensorflow or in keras . I am not fond of using with tf.device(’/device:GPU:2’) - I want the gpu choosing part separate from code writing part.