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

Problem Description

Error when using hiddenlayer to visualize model structure:
RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

Cause Analysis

Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same.
First check whether you are using the GPU for acceleration,

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

Therefore, the model uses GPU, so we need to ensure that the input data is also in GPU form.

solution:

Find the location of the input, add .cuda() after the input, and convert the input from the CPU to the GPU.
error code:

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

Correct code:

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

Finish!