Minor fix in COCO dataset

0332624e · Francisco Massa · 9378febe · 0332624e
Commit 0332624e authored Nov 10, 2016 by Francisco Massa
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

torchvision/datasets/coco.py torchvision/datasets/coco.py +3 -0

No files found.
--- a/torchvision/datasets/coco.py
+++ b/torchvision/datasets/coco.py
@@ -51,6 +51,9 @@ class CocoDetection(data.Dataset):
        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)