Fixes Lightning-Universe#93

Instead of labels, file names were used for making label-to-class map.
ibraheemmmoosa · Feb 9, 2021 · eade7cd · eade7cd
1 parent 253063d
commit eade7cd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flash/vision/classification/data.py b/flash/vision/classification/data.py
@@ -56,7 +56,7 @@ def __init__(
         self.transform = transform
         self.loader = loader
         if self.has_labels:
-            self.label_to_class_mapping = {v: k for k, v in enumerate(list(sorted(list(set(self.fnames)))))}
+            self.label_to_class_mapping = {v: k for k, v in enumerate(list(sorted(list(set(self.labels)))))}
 
     @property
     def has_labels(self) -> bool: