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

HarshRangwala: <p>I am building an image classification model using pytorch.<br>
Here’s my model;</p>
<pre><code class="lang-auto">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)
</code></pre>
<p>After training the model, I get an error when I run this piece of code to display predictions:</p>
<pre><code class="lang-auto">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"))
</code></pre>
<p>Error message is:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------

AttributeError                            Traceback (most recent call last)

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

AttributeError: 'generator' object has no attribute 'next'
</code></pre>
<p>As an alternate, when I try this, I get an error:</p>
<pre><code class="lang-auto">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))
</code></pre>
<p>Error message:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

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

5 frames

&lt;ipython-input-51-38191d042414&gt; in predict_image(img, model)
      1 def predict_image(img, model):
      2     xb = img.unsqueeze(0)
----&gt; 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:
--&gt; 550             result = self.forward(*input, **kwargs)
    551         for hook in self._forward_hooks.values():
    552             hook_result = hook(self, input, result)

&lt;ipython-input-40-f41b37e603a0&gt; in forward(self, x)
     10     self.fc2 = nn.Linear(50, output_size)
     11   def forward(self, x):
---&gt; 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:
--&gt; 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):
--&gt; 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,
--&gt; 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
</code></pre>

harsha_g: <p>Try <code>next(dataiter)</code> instead of <code>dataiter.next()</code></p>

harsha_g: <aside class="quote no-group quote-modified" data-username="HarshRangwala" data-post="1" data-topic="84935">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/harshrangwala/48/25088_2.png" class="avatar"> HarshRangwala:</div>
<blockquote>
<p><code>model(xb).</code></p>
</blockquote>
</aside>
<p>What is next to the <code>.</code>? Probably, your model is on the GPU but the input image is on CPU.</p>

HarshRangwala: <p>Sorry my mistake. There is no <code>.</code> there.</p>

HarshRangwala: <p>When I run that cell I can display the image but cannot get the prediction.</p>

HarshRangwala: <p>This didnt work. Its showing same images.</p>

harsha_g: <p>I can’t think of anything else other than the <code>model</code> and the <code>image</code> being on different devices. Try moving both to the <em>cpu</em> or <em>gpu</em>.</p>

HarshRangwala: <p>yeah. I already moved my model to GPU using:</p>
<pre><code class="lang-auto">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)
</code></pre>
<pre><code class="lang-auto">model = to_device(trafficsignalModel(), device)
</code></pre>
<p>Can you tell me how to do that for image?</p>