Fixed bug

leschultz · Apr 17, 2024 · a9498fd · a9498fd
1 parent bae05b2
commit a9498fd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/madml/plots.py b/src/madml/plots.py
@@ -548,13 +548,13 @@ def __init__(
 
         # Domain prediction columns
         pred_cols = [i for i in self.df.columns if 'Domain Prediction' in i]
-        pred_cols = [i.split(' (')[0] for i in pred_cols]
-        self.pred_cols = pred_cols
 
         if df_confusion is None:
             self.df_confusion = df[self.domains+pred_cols+['splitter']]
+            self.pred_cols = pred_cols
         else:
             self.df_confusion = df_confusion
+            self.pred_cols = [i.split(' (')[0] for i in pred_cols]
 
     def generate(self):