[Pytorch] set device 그래픽카드 지정하기

 

Python Code 내에서 지정

torch.cuda.set_device(1)    # 1번 GPU 사용

 

Python 코드 실행 시 지정

CUDA_VISIBLE_DEVICES=2,3 python my_script.py    # 2, 3번 GPU 사용