I am building an image classification model using pytorch.
Here’s my model;
class trafficsignalModel(ImageClassificationBase):
  
  def __init__(self):
    #super(trafficsignalModel, self).__init__()
    super().__init__()
    self.conv1 = nn.Conv2d(3, 10, kernel_size = 5)
    self.conv2 = nn.Conv2d(10, 20, kernel_size = 5)
    self.conv2_drop = nn.Dropout2d()
    self.fc1 = nn.Linear(500, 50)
    self.fc2 = nn.Linear(50, output_size)
  def forward(self, x):
    x = F.relu(F.max_pool2d(self.conv1(x), 2))
    x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
    x = x.view(-1, 500)
    x = F.relu(self.fc1(x))
    x = F.dropout(x, training=self.training)
    x = self.fc2(x)
    return F.log_softmax(x)

After training the model, I get an error when I run this piece of code to display predictions:
import numpy as np
def im_convert(tensor):
  image = tensor.cpu().clone().detach().numpy()
  image = image.transpose(1, 2, 0)
  image = image * np.array((0.5, 0.5, 0.5)) + np.array((0.5, 0.5, 0.5))
  image = image.clip(0, 1)
  return image


classes = ('0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14',
           '15', '16', '17', '18', '19','20', '21', '22', '23', '24', '25', '26', '27', '28', '29',
           '30', '31', '32', '33', '34', '35', '36', '37', '38', '39','41','42')

dataiter = iter(val_loader)
images, labels = dataiter.next()
output = model(images)
_, preds = torch.max(output, 1)

fig = plt.figure(figsize=(25, 4))

for idx in np.arange(20):
  ax = fig.add_subplot(2, 10, idx+1, xticks=[], yticks=[])
  plt.imshow(im_convert(images[idx]))
  ax.set_title("{} ({})".format(str(classes[preds[idx].item()]), str(classes[labels[idx].item()])), color=("green" if preds[idx]==labels[idx] else "red"))

Error message is:
---------------------------------------------------------------------------

AttributeError                            Traceback (most recent call last)

<ipython-input-48-a549964d8a48> in <module>()
     13 
     14 dataiter = iter(val_loader)
---> 15 images, labels = dataiter.next()
     16 output = model(images)
     17 _, preds = torch.max(output, 1)

AttributeError: 'generator' object has no attribute 'next'

As an alternate, when I try this, I get an error:
def predict_image(img, model):
    xb = img.unsqueeze(0)
    yb = model(xb)
    _, preds  = torch.max(yb, dim=1)
    return preds[0].item()

img, label = test_dataset[11]
plt.imshow(img[0], cmap='gray')
print('Label:', label, ', Predicted:', predict_image(img, model))

Error message:
---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

<ipython-input-51-38191d042414> in <module>()
      7 img, label = test_dataset[11]
      8 plt.imshow(img[0], cmap='gray')
----> 9 print('Label:', label, ', Predicted:', predict_image(img, model))

5 frames

<ipython-input-51-38191d042414> in predict_image(img, model)
      1 def predict_image(img, model):
      2     xb = img.unsqueeze(0)
----> 3     yb = model(xb)
      4     _, preds  = torch.max(yb, dim=1)
      5     return preds[0].item()

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    548             result = self._slow_forward(*input, **kwargs)
    549         else:
--> 550             result = self.forward(*input, **kwargs)
    551         for hook in self._forward_hooks.values():
    552             hook_result = hook(self, input, result)

<ipython-input-40-f41b37e603a0> in forward(self, x)
     10     self.fc2 = nn.Linear(50, output_size)
     11   def forward(self, x):
---> 12     x = F.relu(F.max_pool2d(self.conv1(x), 2))
     13     x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
     14     x = x.view(-1, 500)

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    548             result = self._slow_forward(*input, **kwargs)
    549         else:
--> 550             result = self.forward(*input, **kwargs)
    551         for hook in self._forward_hooks.values():
    552             hook_result = hook(self, input, result)

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/conv.py in forward(self, input)
    347 
    348     def forward(self, input):
--> 349         return self._conv_forward(input, self.weight)
    350 
    351 class Conv3d(_ConvNd):

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight)
    344                             _pair(0), self.dilation, self.groups)
    345         return F.conv2d(input, weight, self.bias, self.stride,
--> 346                         self.padding, self.dilation, self.groups)
    347 
    348     def forward(self, input):

RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _thnn_conv2d_forward