『PyTorch』使用指定 GPU 的方法

CUDA_VISIBLE_DEVICES=id_of_GPU python your_file.py

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "id_of_GPU"

import torch
torch.cuda.set_device(id)

import torch
torch.cuda.current_device()