Merge pull request #18 from rtqichen/master

python 3 fix for cifar test sets.

Merge pull request #18 from rtqichen/master
python 3 fix for cifar test sets.
99ace645 · Soumith Chintala · GitHub · b733ad1a · 486839ed · 99ace645
Commit 99ace645 authored Dec 24, 2016 by Soumith Chintala Committed by GitHub Dec 24, 2016
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

torchvision/datasets/cifar.py torchvision/datasets/cifar.py +4 -1

No files found.
--- a/torchvision/datasets/cifar.py
+++ b/torchvision/datasets/cifar.py
@@ -66,7 +66,10 @@ class CIFAR10(data.Dataset):
            f = self.test_list[0][0]
            file = os.path.join(root, self.base_folder, f)
            fo = open(file, 'rb')
-            entry = pickle.load(fo)
+            if sys.version_info[0] == 2:
+                entry = pickle.load(fo)
+            else:
+                entry = pickle.load(fo, encoding='latin1')
            self.test_data = entry['data']
            if 'labels' in entry:
                self.test_labels = entry['labels']