例子1:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["OMP_NUM_THREADS"] = "1"
model = nn.DataParallel(model).cuda()
device='cuda' # device='cpu'
data = data.to(device)
例子2:
import torch
device = torch.device('cuda') # device = torch.device('cuda:2')
images = images.to(device)