1 files changed, 16 insertions, 1 deletions
diff --git a/utils/configuration.py b/utils/configuration.py
index 455abe8..71584c0 100644
--- a/utils/configuration.py
+++ b/utils/configuration.py
@@ -27,7 +27,7 @@ class DataloaderConfiguration(TypedDict):
     pin_memory: bool
 
 
-class HyperparameterConfiguration(TypedDict):
+class ModelHPConfiguration(TypedDict):
     ae_feature_channels: int
     f_a_c_p_dims: Tuple[int, int, int]
     hpm_scales: Tuple[int, ...]
@@ -41,8 +41,23 @@ class HyperparameterConfiguration(TypedDict):
     tfa_num_parts: int
     embedding_dims: int
     triplet_margin: float
+
+
+class OptimizerHPConfiguration(TypedDict):
     lr: int
     betas: Tuple[float, float]
+    weight_decay: float
+
+
+class SchedulerHPConfiguration(TypedDict):
+    step_size: int
+    gamma: float
+
+
+class HyperparameterConfiguration(TypedDict):
+    model: ModelHPConfiguration
+    optimizer: OptimizerHPConfiguration
+    scheduler: SchedulerHPConfiguration
 
 
 class ModelConfiguration(TypedDict):