Update data argument in predict function

StatMixedML · Jul 20, 2023 · a098e1d · a098e1d
1 parent 5d29ac6
commit a098e1d
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/lightgbmlss/model.py b/lightgbmlss/model.py
@@ -439,7 +439,7 @@ def objective(trial):
         return opt_param.params
 
     def predict(self,
-                test_set: pd.DataFrame,
+                data: pd.DataFrame,
                 pred_type: str = "parameters",
                 n_samples: int = 1000,
                 quantiles: list = [0.1, 0.5, 0.9],
@@ -449,8 +449,8 @@ def predict(self,
 
         Arguments
         ---------
-        test_set : pd.DataFrame
-            Test data.
+        data : pd.DataFrame
+            Data to predict from.
         pred_type : str
             Type of prediction:
             - "samples" draws n_samples from the predicted distribution.
@@ -472,7 +472,7 @@ def predict(self,
 
         # Predict
         predt_df = self.dist.predict_dist(booster=self.booster,
-                                          test_set=test_set,
+                                          data=data,
                                           start_values=self.start_values,
                                           pred_type=pred_type,
                                           n_samples=n_samples,