Also,
Since its in a training/testing loop matplotlib doesnt make much sense.
Rather I was trying to write the image to file.
using torchvision.utils.save_image() for that. However it gives the following error :
pred_imgs = torch.from_numpy(pred_imgs[0])
from torchvision import utils as u 
u.save_image(pred_imgs, './pred_imgs.png')


Traceback (most recent call last):
  File "/home/saleem/anaconda3/envs/khaturia/lib/python3.5/site-packages/PIL/Image.py", line 2460, in fromarray
    mode, rawmode = _fromarray_typemap[typekey]
KeyError: ((1, 1, 150), '|u1')

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/saleem/anaconda3/envs/khaturia/lib/python3.5/site-packages/torchvision-0.2.1-py3.5.egg/torchvision/utils.py", line 103, in save_image
  File "/home/saleem/anaconda3/envs/khaturia/lib/python3.5/site-packages/PIL/Image.py", line 2463, in fromarray
    raise TypeError("Cannot handle this data type")
TypeError: Cannot handle this data type

how to write this image to file ?