[ADD](handlers) Train cfg to main cfg base model

wbenbihi · Aug 23, 2022 · 62498ab · 62498ab
1 parent 3ca77bc
commit 62498ab
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/hourglass_tensorflow/types/config/__init__.py b/hourglass_tensorflow/types/config/__init__.py
@@ -11,6 +11,7 @@
 
 from hourglass_tensorflow.types.config.data import HTFDataConfig
 from hourglass_tensorflow.types.config.model import HTFModelConfig
+from hourglass_tensorflow.types.config.train import HTFTrainConfig
 from hourglass_tensorflow.types.config.fields import HTFConfigField
 from hourglass_tensorflow.utils.object_logger import ObjectLogger
 from hourglass_tensorflow.types.config.dataset import HTFDatasetConfig
@@ -23,11 +24,12 @@ class HTFConfigMode(enum.Enum):
 
 
 class HTFConfig(HTFConfigField):
-    version: Optional[Union[str, int]]
     mode: HTFConfigMode = HTFConfigMode.TRAIN
+    version: Optional[Union[str, int]]
     data: Optional[HTFDataConfig]
     dataset: Optional[HTFDatasetConfig]
     model: Optional[HTFModelConfig]
+    train: Optional[HTFTrainConfig]
 
 
 # region ConfigParser

diff --git a/hourglass_tensorflow/types/config/train.py b/hourglass_tensorflow/types/config/train.py
@@ -15,6 +15,6 @@
 class HTFTrainConfig(HTFConfigField):
     object: Optional[HTFObjectReference] = Field(
         default_factory=HTFObjectReference(
-            source="hourglass_tensorflow.handlers.model.HTFModelHandler"
+            source="hourglass_tensorflow.handlers.train.HTFTrainHandler"
         )
     )