From 0d463c119355cd37f3ea06cd5868938fb76cb0a6 Mon Sep 17 00:00:00 2001
From: Mike Vink <mike1994vink@gmail.com>
Date: Fri, 25 Sep 2020 08:29:30 +0200
Subject: minor things

---
 tree.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

(limited to 'tree.py')

diff --git a/tree.py b/tree.py
index e971bd4..3a073ac 100644
--- a/tree.py
+++ b/tree.py
@@ -1,10 +1,10 @@
 import numpy as np
-import cProfile
-import pstats
+# import cProfile
+# import pstats
 # import tqdm
 
 # from tqdm import trange
-from pstats import SortKey
+# from pstats import SortKey
 from sklearn import metrics
 
 # age,married,house,income,gender,class
@@ -252,7 +252,7 @@ def update_mask(mask, current_mask):
     equal to the total number of rows in dataset x.
     """
     copy = np.array(current_mask, copy=True)
-    copy[np.where(current_mask)] = mask
+    copy[current_mask == True] = mask
     return copy
 
 
-- 
cgit v1.2.3