DS_Store fix in find_classes (#79)

pytorch · Feb 27, 2017 · c0a6cfe · c0a6cfe
1 parent 39a50de
commit c0a6cfe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchvision/datasets/folder.py b/torchvision/datasets/folder.py
@@ -15,7 +15,7 @@ def is_image_file(filename):
 
 
 def find_classes(dir):
-    classes = os.listdir(dir)
+    classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
     classes.sort()
     class_to_idx = {classes[i]: i for i in range(len(classes))}
     return classes, class_to_idx