fix(Metalearner): Fixed postprocessing of Decision Tree prediction

frankkramer-lab · May 22, 2022 · 72794df · 72794df
1 parent fedac2d
commit 72794df
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/aucmedi/ensemble/metalearner/decision_tree.py b/aucmedi/ensemble/metalearner/decision_tree.py
@@ -62,6 +62,9 @@ def train(self, x, y):
     def predict(self, data):
         # Compute prediction probabilities via fitted model
         pred = self.model.predict_proba(data)
+        # Postprocess decision tree predictions
+        pred = np.asarray(pred)
+        pred = np.swapaxes(pred[:,:,1], 0, 1)
         # Return results as NumPy array
         return pred