开头加一段这个代码

作用：首先检查CUDA是否可用，如果可用，则将默认张量类型设置为CUDA的浮点张量，并将设备设置为cuda，否则设置为cpu

if torch.cuda.is_available():
  torch.set_default_tensor_type(torch.cuda.FloatTensor)
  print("using cuda:", torch.cuda.get_device_name(0))

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

训练模型时，需将模型转移到gpu上

模型名.to(device)

创建pytorch张量

 torch.cuda.FloatTensor([1.0])

注：GPU tensor要先转到CPU tensor，才能转换为numpy数组

re = output.detach().cpu().numpy()