Merge pull request #1 from fmassa/coco_tfix

Minor fix in COCO dataset
pytorch · Nov 10, 2016 · e37323d · e37323d
2 parents 9378feb + 0332624
commit e37323d
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/torchvision/datasets/coco.py b/torchvision/datasets/coco.py
@@ -51,6 +51,9 @@ def __getitem__(self, index):
         path = coco.loadImgs(img_id)[0]['file_name']
 
         img = Image.open(os.path.join(self.root, path)).convert('RGB')
+        if self.transform is not None:
+            img = self.transform(img)
+
         if self.target_transform is not None:
             target = self.target_transform(target)