Merge pull request #103 from tlapusan/fix_bug_issue_101

Fix bug in case xgboost model has only one tree

dtreeviz/models/xgb_decision_tree.py

@@ -93,7 +93,11 @@ def get_node_samples(self):
             return self.node_to_samples
 
         prediction_leaves = self.booster.predict(xgb.DMatrix(self.x_data, feature_names=self.feature_names),
-                                                 pred_leaf=True)[:, self.tree_index]
+                                                 pred_leaf=True)
+
+        if len(prediction_leaves.shape) > 1:
+            prediction_leaves = prediction_leaves[:, self.tree_index]
+
         node_to_samples = defaultdict(list)
         for sample_i, prediction_leaf in enumerate(prediction_leaves):
             prediction_path = self._get_leaf_prediction_path(prediction_leaf)