Search code examples
jupyter-notebookpytorchgpu

How to make Jupyter Notebook run on one specified gpu of multi-gpus

I have some PyTorch code in one Jupyter Notebook which needs to run on one specified GPU (that is, not 'GPU 0') since others already work on 'GPU 0'. I have tried this:

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

But this does not work. How do I run the code on one specified GPU successfully?

Solution

  • You can set the environment variable for CUDA_VISIBLE_DEVICES at the beginning of the same Jupyter Notebook cell that has the code that puts the model on the device.

    Set the environment variable with:

    !export CUDA_VISIBLE_DEVICES=4,5,6,7

    And not in Python, neither

    os.environ["CUDA_VISIBLE_DEVICES"]="4,5,6,7"

    nor:

    %set_env CUDA_VISIBLE_DEVICES=4,5,6,7

    There are ways to make this setting a default: .bashrc, kernel.json, Jupyter Service settings (all untested), see Stack Overflow How to set env variable in Jupyter notebook.

    Proof

    Here is a code proof that setting the CUDA_VISIBLE_DEVICES env var in Python makes already working code that runs !export ... in the beginning fail again, see How to Solve "CUDA: Invalid Device Ordinal" Error in PyTorch Single-GPU Inference on a Model Trained with DataParallel?.