Update type hint

BirkhoffG · Apr 4, 2024 · 1cb16be · 1cb16be
1 parent 72fc25c
commit 1cb16be
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/nbs/01_data.ipynb b/nbs/01_data.ipynb
@@ -550,7 +550,7 @@
     "    \n",
     "    def set_transformations(\n",
     "        self, \n",
-    "        feature_names_to_transformation: Dict[str, Union[str, Dict, Transformation]], # Dict[feature_name, Transformation]\n",
+    "        feature_names_to_transformation: Dict[str, Union[str, Dict, BaseTransformation]], # Dict[feature_name, Transformation]\n",
     "    ) -> DataModule:\n",
     "        \"\"\"Reset transformations for features.\"\"\"\n",
     "\n",

diff --git a/relax/data_module.py b/relax/data_module.py
@@ -300,7 +300,7 @@ def __getitem__(self, name: str):
 
     def set_transformations(
         self, 
-        feature_names_to_transformation: Dict[str, Union[str, Dict, Transformation]], # Dict[feature_name, Transformation]
+        feature_names_to_transformation: Dict[str, Union[str, Dict, BaseTransformation]], # Dict[feature_name, Transformation]
     ) -> DataModule:
         """Reset transformations for features."""