diff --git a/flaml/tune/searcher/suggestion.py b/flaml/tune/searcher/suggestion.py index 8380fc0fda..b9bca38127 100644 --- a/flaml/tune/searcher/suggestion.py +++ b/flaml/tune/searcher/suggestion.py @@ -561,7 +561,15 @@ def __init__( self._space = space self._points_to_evaluate = points_to_evaluate or [] - self._evaluated_rewards = evaluated_rewards + # rewards should be a list of floats, not a dict + # After Optuna > 3.5.0, there is a check for NaN in the list "any(math.isnan(x) for x in self._values)" + # which will raise an error when encountering a dict + if evaluated_rewards is not None: + self._evaluated_rewards = [ + list(item.values())[0] if isinstance(item, dict) else item for item in evaluated_rewards + ] + else: + self._evaluated_rewards = evaluated_rewards self._study_name = "optuna" # Fixed study name for in-memory storage diff --git a/setup.py b/setup.py index 1c6d045bfe..7bb35716bb 100644 --- a/setup.py +++ b/setup.py @@ -65,7 +65,7 @@ "catboost>=0.26,<1.2; python_version<'3.11'", "catboost>=0.26; python_version>='3.11'", "rgf-python", - "optuna==2.8.0", + "optuna>=2.8.0,<=3.6.1", "openml", "statsmodels>=0.12.2", "psutil==5.8.0", @@ -94,7 +94,7 @@ ], "catboost": ["catboost>=0.26"], "blendsearch": [ - "optuna==2.8.0", + "optuna>=2.8.0,<=3.6.1", "packaging", ], "ray": [ @@ -153,7 +153,7 @@ ], "synapse": [ "joblibspark>=0.5.0", - "optuna==2.8.0", + "optuna>=2.8.0,<=3.6.1", "pyspark>=3.2.0", ], "autozero": ["scikit-learn", "pandas", "packaging"],