diff --git a/image_ref/main.py b/image_ref/main.py
index 3d983bc765b526e5298f9d7b5aaaa26157f664b2..12a696e6222cc3fb470b25081be23b9678df130c 100644
--- a/image_ref/main.py
+++ b/image_ref/main.py
@@ -169,7 +169,7 @@ def make_prediction_duo(model, data, f_name, f_name2):
     # constant for classes
 
     # Build confusion matrix
-    classes = data.dataset.dataset.classes
+    classes = data.dataset.classes
     cf_matrix = confusion_matrix(y_true, y_pred)
     confidence_matrix = np.zeros((n_class,n_class))
     for i in range(n_class):