diff --git a/aepsych/generators/completion_criterion/min_total_tells.py b/aepsych/generators/completion_criterion/min_total_tells.py index 25c4887f4..1a8bfff9f 100644 --- a/aepsych/generators/completion_criterion/min_total_tells.py +++ b/aepsych/generators/completion_criterion/min_total_tells.py @@ -9,12 +9,15 @@ from aepsych.config import Config, ConfigurableMixin from ax.core.base_trial import TrialStatus -from ax.modelbridge.completion_criterion import MinimumTrialsInStatus +from ax.modelbridge.transition_criterion import MinTrials -class MinTotalTells(MinimumTrialsInStatus, ConfigurableMixin): +class MinTotalTells(MinTrials, ConfigurableMixin): @classmethod def get_config_options(cls, config: Config, name: str) -> Dict[str, Any]: min_total_tells = config.getint(name, "min_total_tells", fallback=1) - options = {"status": TrialStatus.COMPLETED, "threshold": min_total_tells} + options = { + "only_in_statuses": [TrialStatus.COMPLETED], + "threshold": min_total_tells, + } return options