There is one mistake I noticed: after the training when you are passing images through your network, you should also transfer your images to the GPU because your net is defined on GPU. For instance it should be
images, labels = images.to(device), labels.to(device)
anytime before you use net(images).