diff --git a/src/flux_1_schnell/config/model_config.py b/src/flux_1_schnell/config/model_config.py
index 0f70952..d236cc2 100644
--- a/src/flux_1_schnell/config/model_config.py
+++ b/src/flux_1_schnell/config/model_config.py
@@ -2,12 +2,19 @@ from enum import Enum
 
 
 class ModelConfig(Enum):
-    FLUX1_DEV = ("black-forest-labs/FLUX.1-dev", "dev", 512)
-    FLUX1_SCHNELL = ("black-forest-labs/FLUX.1-schnell", "schnell", 256)
+    FLUX1_DEV = ("black-forest-labs/FLUX.1-dev", "dev", 1000, 512)
+    FLUX1_SCHNELL = ("black-forest-labs/FLUX.1-schnell", "schnell", 1000, 256)
 
-    def __init__(self, model_name: str, alias: str, max_sequence_length: int):
+    def __init__(
+            self,
+            model_name: str,
+            alias: str,
+            num_train_steps: int,
+            max_sequence_length: int,
+    ):
         self.alias = alias
         self.model_name = model_name
+        self.num_train_steps = num_train_steps
         self.max_sequence_length = max_sequence_length
 
     @staticmethod
diff --git a/src/flux_1_schnell/config/runtime_config.py b/src/flux_1_schnell/config/runtime_config.py
index 3807ec1..e4b3937 100644
--- a/src/flux_1_schnell/config/runtime_config.py
+++ b/src/flux_1_schnell/config/runtime_config.py
@@ -9,7 +9,7 @@ class RuntimeConfig:
 
     def __init__(self, config: Config, model_config: ModelConfig):
         self.config = config
-        self.num_train_steps = 1000
+        self.model_config = model_config
         self.sigmas = self._create_sigmas(config, model_config)
 
     @property
@@ -32,6 +32,10 @@ class RuntimeConfig:
     def precision(self):
         return self.config.precision
 
+    @property
+    def num_train_steps(self):
+        return self.model_config.num_train_steps
+
     @staticmethod
     def _create_sigmas(config, model):
         sigmas = RuntimeConfig._create_sigmas_values(config.num_inference_steps)