How do I solve the generator object error and/or convert device type from cpu to cuda?

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

Try next(dataiter) instead of dataiter.next()

What is next to the .? Probably, your model is on the GPU but the input image is on CPU.

Sorry my mistake. There is no . there.

When I run that cell I can display the image but cannot get the prediction.

This didnt work. Its showing same images.

I can’t think of anything else other than the model and the image being on different devices. Try moving both to the cpu or gpu.

yeah. I already moved my model to GPU using:

def to_device(data, device):
    """Move tensor(s) to chosen device"""
    if isinstance(data, (list,tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device, dtype = torch.float, non_blocking=True)

class DeviceDataLoader():
    """Wrap a dataloader to move data to a device"""
    def __init__(self, dl, device):
        self.dl = dl
        self.device = device
        
    def __iter__(self):
        """Yield a batch of data after moving it to device"""
        for b in self.dl: 
            yield to_device(b, self.device)

    def __len__(self):
        """Number of batches"""
        return len(self.dl)
model = to_device(trafficsignalModel(), device)

Can you tell me how to do that for image?