Catch AttributeError for newer versions of TPOT

openml · Jun 17, 2023 · be6d6bb · be6d6bb
1 parent e1b1cd7
commit be6d6bb
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/frameworks/TPOT/exec.py b/frameworks/TPOT/exec.py
@@ -71,7 +71,7 @@ def infer(data):
         if is_classification:
             try:
                 return tpot.predict_proba(data)
-            except RuntimeError:
+            except (RuntimeError, AttributeError):
                 return tpot.predict(data)
         return tpot.predict(data)
 
@@ -94,8 +94,9 @@ def infer(data):
 
     try:
         probabilities = tpot.predict_proba(X_test) if is_classification else None
-    except RuntimeError:
-        # TPOT throws a RuntimeError if the optimized pipeline does not support `predict_proba`.
+    except (RuntimeError, AttributeError):
+        # TPOT throws a RuntimeError or AttributeError if the optimized pipeline
+        # does not support `predict_proba` (which one depends on the version).
         probabilities = "predictions"  # encoding is handled by caller in `__init__.py`
 
     save_artifacts(tpot, config)