I tried adding .cuda() to fix this:
images = images.unsqueeze(1).type(torch.FloatTensor).cuda()
outputs = model(images)

But when I do this it brings back the original error: “RuntimeError: Expected 4-dimensional input for 4-dimensional weight [16, 1, 5, 5], but got 3-dimensional input of size [5, 28, 28] instead”, I don’t know why .cuda() would affect the dimensions.
EDIT: I fixed this problem, it’s because I wasn’t doing the update to the dimensions correctly further down on the accuracy calculation. The code runs correctly now! I’m just getting “nan” for the loss and accuracy of “4.615”, which does not make much sense… I’ll continue looking over my code. Thanks for your help.
EDIT: I fixed the nan problem by reducing the learning rate. All good now!