Update test_common.py

uxlfoundation · Jan 13, 2025 · 6eaec9a · 6eaec9a
1 parent 9de1f51
commit 6eaec9a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/sklearnex/tests/test_common.py b/sklearnex/tests/test_common.py
@@ -247,7 +247,10 @@ def sklearnex_trace(estimator_name, method_name):
         of trace._modname.
     """
     # get estimator
-    estimator = (SPECIAL_INSTANCES | PATCHED_MODELS)[estimator_name]
+    try:
+        est = PATCHED_MODELS[estimator_name]()
+    except KeyError:
+        est = SPECIAL_INSTANCES[estimator_name]
 
     # get dataset
     X, y = gen_dataset(est)[0]