Added ability to set timesteps to linear for flowmatching schedule

2024-08-11 13:06:08 -06:00
parent f8f0657b68
commit a6aa4b2c7d
3 changed files with 19 additions and 11 deletions
--- a/toolkit/config_modules.py
+++ b/toolkit/config_modules.py
@@ -356,6 +356,7 @@ class TrainConfig:
        # adds an additional loss to the network to encourage it output a normalized standard deviation
        self.target_norm_std = kwargs.get('target_norm_std', None)
        self.target_norm_std_value = kwargs.get('target_norm_std_value', 1.0)
+        self.linear_timesteps = kwargs.get('linear_timesteps', False)


 class ModelConfig: