Check for superset instead of set equality.

Quantco · Jun 18, 2024 · c7d6ef2 · c7d6ef2
1 parent f5d0aa3
commit c7d6ef2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/metalearners/metalearner.py b/metalearners/metalearner.py
@@ -93,8 +93,8 @@ def _get_result_skeleton(
 
 
 def _initialize_model_dict(argument, expected_names: Collection[str]) -> dict:
-    if isinstance(argument, dict) and set(argument.keys()) == set(expected_names):
-        return argument
+    if isinstance(argument, dict) and set(argument.keys()) >= set(expected_names):
+        return {key: argument[key] for key in expected_names}
     return {name: argument for name in expected_names}