From dee702260c2e0a09f4515de1cd9f69a22ee705c9 Mon Sep 17 00:00:00 2001 From: Vasilev Dmitrii <70897456+D1MK4real@users.noreply.github.com> Date: Thu, 10 Aug 2023 15:13:27 +0300 Subject: [PATCH] float to int (#111) --- lightautoml/ml_algo/dl_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lightautoml/ml_algo/dl_model.py b/lightautoml/ml_algo/dl_model.py index ac1dae5d..0460fa05 100644 --- a/lightautoml/ml_algo/dl_model.py +++ b/lightautoml/ml_algo/dl_model.py @@ -340,7 +340,7 @@ def _init_params_on_input(self, train_valid_iterator) -> dict: target = train_valid_iterator.train.target if params["n_out"] is None: - new_params["n_out"] = 1 if task_name != "multiclass" else np.max(target) + 1 + new_params["n_out"] = 1 if task_name != "multiclass" else (np.max(target) + 1).astype(int) new_params["n_out"] = target.shape[1] if task_name in ["multi:reg", "multilabel"] else new_params["n_out"] cat_dims = []