This is the error I am getting when I run
visualize_model(net)

Error message
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-43-7ec848dbaa28> in <module>()
      1 
----> 2 visualize_model(net)

<ipython-input-41-3b3aaf3229b1> in visualize_model(net, num_images)
     18                 ax.axis('off')
     19                 ax.set_title('predicted: {}'.format(class_names[preds[j]]))
---> 20                 imshow(inputs.data[j])
     21 
     22                 if images_so_far == num_images:

<ipython-input-40-315a2e968437> in imshow(img)
      1 def imshow(img):
      2     #img = img / 2 + 0.5
----> 3     plt.imshow(np.transpose(img[0].numpy(), (1, 2, 0)))
      4     plt.show()
      5 

C:\ProgramData\Anaconda3\lib\site-packages\numpy\core\fromnumeric.py in transpose(a, axes)
    596 
    597     """
--> 598     return _wrapfunc(a, 'transpose', axes)
    599 
    600 

C:\ProgramData\Anaconda3\lib\site-packages\numpy\core\fromnumeric.py in _wrapfunc(obj, method, *args, **kwds)
     49 def _wrapfunc(obj, method, *args, **kwds):
     50     try:
---> 51         return getattr(obj, method)(*args, **kwds)
     52 
     53     # An AttributeError occurs if the object does not have

ValueError: axes don't match array