# 选择设备，选择使用GPU(cuda)训练

# first way
if torch.cuda.is_available():
    device = torch.device("cuda:0")
else:
    device = torch.device("cpu")


# second way
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")