解决办法Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

问题描述

在使用hiddenlayer进行模型结构可视化时报错：
RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

原因分析

输入类型（torch.FloatTensor）和权重类型（torch.cuda.FloatTensor）应该相同。
首先检查自己是否使用了GPU进行加速，

if args.gpu:
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
model = Net(args).cuda()

因此模型使用了GPU，所以我们要保证输入的数据也是GPU形式的。

解决方案：

找到输入的位置，在输入后面加上.cuda()即可，将输入由CPU转换成GPU即可。
错误代码：

 vis_graph = h.build_graph(model, torch.zeros([1 ,3, 32,32])) 

正确代码：

vis_graph = h.build_graph(model, torch.zeros([1 ,3, 32,32]).cuda()) 

完成！