使用torchsummary出现报错RuntimeError: cuda runtime error (30) : unknown error at ..\aten\src\THC\THCGenera

使用torchsummary出现报错RuntimeError: cuda runtime error (30) : unknown error at …\aten\src\THC\THCGenera
在语句的前面加上

torch.cuda.current_device()

就可以了

import torch
from torchvision import models
from torchsummary import summary
from models import SegNet
print(torch.cuda.current_device())  #这句必须加
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vgg = models.vgg16().to('cuda')
model =SegNet(2).cuda()
summary(model, (3,640, 480))
summary(vgg, (3,224, 224))
#from torchsummary import summary
 #summary(your_model, input_size=(channels, H, W))