fix gpu test

dask-contrib · Jan 31, 2023 · 98c42d5 · 98c42d5
1 parent 4717bde
commit 98c42d5
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 19 deletions.
diff --git a/tests/integration/test_model.py b/tests/integration/test_model.py
@@ -62,20 +62,20 @@ def test_training_and_prediction(c, gpu_client):
         )
         check_trained_model(c)
 
-    c.sql(
-        f"""
-        CREATE OR REPLACE MODEL my_model WITH (
-            model_class = 'LogisticRegression',
-            wrap_predict = True,
-            wrap_fit = False,
-            target_column = 'target'
-        ) AS (
-            SELECT x, y, x*y > 0 AS target
-            FROM {timeseries}
+        c.sql(
+            f"""
+            CREATE OR REPLACE MODEL my_model WITH (
+                model_class = 'LogisticRegression',
+                wrap_predict = True,
+                wrap_fit = False,
+                target_column = 'target'
+            ) AS (
+                SELECT x, y, x*y > 0 AS target
+                FROM {timeseries}
+            )
+        """
         )
-    """
-    )
-    check_trained_model(c, df_name=timeseries)
+        check_trained_model(c, df_name=timeseries)
 
     c.sql(
         f"""

diff --git a/tests/unit/test_ml_utils.py b/tests/unit/test_ml_utils.py
@@ -44,12 +44,7 @@ def test_ml_class_mappings(gpu):
         if not ("XGB" in key and xgboost is None) and not (
             "LGBM" in key and lightgbm is None
         ):
-            if gpu and key == "LogisticRegression":
-                # dask-glm >= 0.2.1.dev needed to use multi-GPU logistic regression
-                with pytest.raises(ImportError):
-                    import_class(classes_dict[key])
-            else:
-                import_class(classes_dict[key])
+            import_class(classes_dict[key])
 
 
 def _check_axis_partitioning(chunks, n_features):