Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

在使用torchsummary.summary时
遇到Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same问题.

错误内容大概就是指输入类型是CPU(torch.FloatTensor),而参数类型是GPU(torch.cuda.FloatTensor)。

两种用法:

model = torchvision.models.alexnet()
    if vis_model:
        from torchsummary import summary
        summary(model, input_size=(3, 224, 224), device='cpu')  # 输入模型与图片的size,输出网络结构的可视化
model.to(device)

因为直接输入的model是CPU类型,所以summary里面必须写device=‘cpu’

model = torchvision.models.alexnet()
model.to(device)
    if vis_model:
        from torchsummary import summary
        summary(model, input_size=(3, 224, 224))  # 输入模型与图片的size,输出网络结构的可视化

下面这种错误:

model = torchvision.models.alexnet()

    if vis_model:
        from torchsummary import summary
        summary(model, input_size=(3, 224, 224))  # 输入模型与图片的size,输出网络结构的可视
        
model.to(device)

参考

你可能感兴趣的:(#,Pytorch,深度学习,可视化)