Merge pull request #10 from pytorch/cifarfix

making cifar data loader also return PIL Image

Author: soumith

test/sanity_checks.ipynb

@@ -77,6 +77,10 @@ def __getitem__(self, index):
             img, target = self.train_data[index], self.train_labels[index]
         else:
             img, target = self.test_data[index], self.test_labels[index]
+            
+        # doing this so that it is consistent with all other datasets
+        # to return a PIL Image
+        img = Image.fromarray(np.transpose(img, (1,2,0)))
 
         if self.transform is not None:
             img = self.transform(img)

torchvision/datasets/cifar.py

@@ -1,10 +1,20 @@
+import torch
+import math
 
 def make_grid(tensor, nrow=8, padding=2):
     """
     Given a 4D mini-batch Tensor of shape (B x C x H x W),
+    or a list of images all of the same size,
     makes a grid of images
     """
-    import math
+    tensorlist = None
+    if isinstance(tensor, list):
+        tensorlist = tensor
+        numImages = len(tensorlist)
+        size = torch.Size(torch.Size([long(numImages)]) + tensorlist[0].size())
+        tensor = tensorlist[0].new(size)
+        for i in range(numImages):
+            tensor[i].copy_(tensorlist[i])
     if tensor.dim() == 3: # single image
         return tensor
     # make the mini-batch of images into a grid