allow trained model to get num_iters and objective

rapidsai · Sep 28, 2023 · e72f4e9 · e72f4e9
1 parent 8bed6b0
commit e72f4e9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/python/cuml/linear_model/logistic_regression_mg.pyx b/python/cuml/linear_model/logistic_regression_mg.pyx
@@ -209,6 +209,10 @@ class LogisticRegressionMG(MGFitMixin, LogisticRegression):
                 <float*> &objective32,
                 <int*> &num_iters)
 
+            self.solver_model.objective = objective32
+
+        self.solver_model.num_iters = num_iters
+
         self.solver_model._calc_intercept()
 
         self.handle.sync()
diff --git a/python/cuml/tests/dask/test_dask_logistic_regression.py b/python/cuml/tests/dask/test_dask_logistic_regression.py
@@ -20,8 +20,6 @@
 from sklearn.datasets import make_classification
 from sklearn.linear_model import LogisticRegression as skLR
 from cuml.internals.safe_imports import cpu_only_import
-from hypothesis import given
-from hypothesis import strategies as st
 
 pd = cpu_only_import("pandas")
 np = cpu_only_import("numpy")