CUDA_VISIBLE_DEVICESの設定方法

1. bashスクリプトでCUDA_VISIBLE_DEVICESを設定する。

export CUDA_VISIBLE_DEVICES=0,1,2,3
# Be careful not to use spaces in bash scripts, especially not after symbols
source ~/anaconda3/bin/activate torch
# Use the source command to read and execute activate to activate the anaconda virtual environment
python -c 'import torch; print(torch.__version__); print(torch.cuda.device_count())'
# Display pytorch version in python command mode

import torch
import os
os.environ['CUDA_VISIBLE_DEVICES']='0,1,2'
print(torch.cuda.device_count())

2. os.environでCUDA_VISIBLE_DEVICESを設定する。

import torch
import os
os.environ['CUDA_VISIBLE_DEVICES']='0,1,2'
print(torch.cuda.device_count())