Class weights being initialized with only training partition, not val…

…idation
cavalab · Nov 8, 2024 · 2451619 · 2451619
1 parent b0b485d
commit 2451619
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/pybrush/BrushEstimator.py b/pybrush/BrushEstimator.py
@@ -72,7 +72,7 @@ def fit(self, X, y):
         self.train_.set_batch_size(self.batch_size) # TODO: update batch indexes at the beggining of every generation
         self.validation_ = self.data_.get_validation_data()
 
-        self.parameters_ = self._wrap_parameters(y)
+        self.parameters_ = self._wrap_parameters(self.train_.y)
 
         self.search_space_ = SearchSpace(self.data_,
                                          self.parameters_.functions,

diff --git a/pybrush/DeapEstimator.py b/pybrush/DeapEstimator.py
@@ -150,7 +150,7 @@ def fit(self, X, y):
 
         self.validation_ = self.data_.get_validation_data()
 
-        self.parameters_ = self._wrap_parameters(y)
+        self.parameters_ = self._wrap_parameters(self.train_.y)
         self.search_space_ = SearchSpace(self.data_, self.parameters_.functions, self.weights_init)
 
         if self.mode == "classification":