Pytorch--检查是否正常运行(Code)

Code (refer to this link)

import torch
flag = torch.cuda.is_available() # check if the gpu is avaliable
print(flag)
ngpu= 1
# Decide which device we want to run on
device = torch.device(“cuda:0” if (torch.cuda.is_available() and ngpu > 0) else “cpu”)
print(device)
print(torch.cuda.get_device_name(0))
print(torch.rand(3,3).cuda())

Note:

• 建议检查一下Gpu上的显存情况，一般运行后，会占用200m左右显存；