1

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

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

But this is not work! So how to make the code run on one specified gpu successfully?

bukejie
  • 11
  • 3

1 Answers1

0

You may launch Jupyter Notebook with, say, CUDA_VISIBLE_DEVICES=1. This usually what I do on cluster, because PyTorch doc recommends setting CUDA_VISIBLE_DEVICES compared to torch functions like torch.cuda.set_device(device):

$ CUDA_VISIBLE_DEVICES=1 jupyter notebook &

You can also check what device is available in your notebook using torch.cuda.is_available()

Mark
  • 336
  • 1
  • 7