cifar10/100 only load train/test and not both

torchvision/datasets/cifar.py

@@ -42,6 +42,7 @@ class CIFAR10(data.Dataset):
                                + ' You can use download=True to download it')
                 
         # now load the picked numpy arrays
+        if self.train:
             self.train_data = []
             self.train_labels = []
             for fentry in self.train_list:
@@ -57,7 +58,8 @@ class CIFAR10(data.Dataset):
                 fo.close()
 
             self.train_data = np.concatenate(self.train_data)
-
+            self.train_data = self.train_data.reshape((50000, 3, 32, 32))
+        else:
             f = self.test_list[0][0]
             file = os.path.join(root, self.base_folder, f)
             fo = open(file, 'rb')
@@ -68,8 +70,6 @@ class CIFAR10(data.Dataset):
             else:
                 self.test_labels = entry['fine_labels']
             fo.close()
-
-        self.train_data = self.train_data.reshape((50000, 3, 32, 32))
             self.test_data = self.test_data.reshape((10000, 3, 32, 32))
 
     def __getitem__(self, index):