device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')#申明用GPU

如果有多个GPU，你可以选择具体的GPU进行使用。

import os 

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]='1'

1. CPU tensor转GPU tensor：

cpu_imgs.cuda()

2. GPU tensor 转CPU tensor：

gpu_imgs.cpu()

3. numpy转为CPU tensor：

torch.from_numpy( imgs )

4.CPU tensor转为numpy数据：

cpu_imgs.numpy()

5. note：GPU tensor不能直接转为numpy数组，必须先转到CPU tensor。

gpu tensor.cpu().numpy()

6. 如果tensor是标量的话，可以直接使用 item() 函数（只能是标量）将值取出来：

print (loss_output.item())

检验模块

import torch
# 以下代码只有在PyTorch GPU版本上才会执⾏行行
if torch.cuda.is_available():
    device = torch.device("cuda") # GPU
x = torch.ones(2,2)
x = x.to(device)
#x = x.cuda()
print(x.device)