-torchvision.utils.save_image(train_set.view(train_set.size(0), 1, train_set.size(1), train_set.size(2)), 'example.png')
+x = x.view(x.size(0), 1, x.size(1), x.size(2))
+
+x.div_(255)
+
+torchvision.utils.save_image(x, 'example.png')
+
+print('Wrote example.png')