cuda用法

例子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)