Fix dataset mypy error

secondmind-labs · Nov 23, 2023 · 25da01b · 25da01b
1 parent e5cccc8
commit 25da01b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/trieste/ask_tell_optimization.py b/trieste/ask_tell_optimization.py
@@ -192,7 +192,10 @@ def __init__(
         if not isinstance(models, Mapping):
             models = {OBJECTIVE: models}
 
+        self._filtered_datasets = datasets
+
         # reassure the type checker that everything is tagged
+        datasets = cast(Dict[Tag, Dataset], datasets)
         models = cast(Dict[Tag, TrainableProbabilisticModelType], models)
 
         # Get set of dataset and model keys, ignoring any local tag index. That is, only the
@@ -206,7 +209,6 @@ def __init__(
             )
 
         self._datasets = datasets
-        self._filtered_datasets = datasets
         self._models = models
 
         self._query_plot_dfs: dict[int, pd.DataFrame] = {}