Try fix lgbm

mars-project · Oct 3, 2023 · 624fb5a · 624fb5a
1 parent 2750b48
commit 624fb5a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/mars/learn/contrib/lightgbm/_predict.py b/mars/learn/contrib/lightgbm/_predict.py
@@ -78,7 +78,7 @@ def __call__(self):
         elif hasattr(self.model, "classes_"):
             dtype = np.array(self.model.classes_).dtype
         else:
-            dtype = getattr(self.model, "out_dtype_", np.dtype("float"))
+            dtype = getattr(self.model, "out_dtype_", [np.dtype("float")])[0]
 
         if self.output_types[0] == OutputType.tensor:
             # tensor

diff --git a/mars/learn/contrib/lightgbm/_train.py b/mars/learn/contrib/lightgbm/_train.py
@@ -406,11 +406,11 @@ def execute(cls, ctx, op: "LGBMTrain"):
                 op.model_type == LGBMModelType.RANKER
                 or op.model_type == LGBMModelType.REGRESSOR
             ):
-                model.set_params(out_dtype_=np.dtype("float"))
+                model.set_params(out_dtype_=[np.dtype("float")])
             elif hasattr(label_val, "dtype"):
-                model.set_params(out_dtype_=label_val.dtype)
+                model.set_params(out_dtype_=[label_val.dtype])
             else:
-                model.set_params(out_dtype_=label_val.dtypes[0])
+                model.set_params(out_dtype_=[label_val.dtypes[0]])
 
             ctx[op.outputs[0].key] = pickle.dumps(model)
         finally: