Not so good visualization results

[Elaineok]

@karanchahal
I have two problems.
1.During training the model（ train_saliency.py），loss becomes NAN after epoch=4.
2.I run the visualize.py but the mask not so good. What's the problem here?
imshow(torchvision.utils.make_grid(masks.cpu().data))

Thanks for your response.