update

examples/nx_clean_unet/yaml/config.yaml

@@ -22,5 +22,9 @@ discriminator_in_channel: 2
 
 compress_factor: 0.3
 
+batch_size: 4
+learning_rate: 0.0005
+adam_b1: 0.8
+adam_b2: 0.99
+lr_decay: 0.99
 seed: 1234
-batch_size: 3

toolbox/torchaudio/models/nx_clean_unet/configuration_nx_clean_unet.py

@@ -28,8 +28,13 @@ class NXCleanUNetConfig(PretrainedConfig):
 
                  compress_factor: float = 0.3,
 
+                 batch_size: int = 4,
+                 learning_rate: float = 0.0005,
+                 adam_b1: float = 0.8,
+                 adam_b2: float = 0.99,
+                 lr_decay: float = 0.99,
                  seed: int = 1234,
-                 batch_size: int = 3,
+
                  **kwargs
                  ):
         super(NXCleanUNetConfig, self).__init__(**kwargs)
@@ -53,8 +58,12 @@ class NXCleanUNetConfig(PretrainedConfig):
 
         self.compress_factor = compress_factor
 
-        self.seed = seed
         self.batch_size = batch_size
+        self.learning_rate = learning_rate
+        self.adam_b1 = adam_b1
+        self.adam_b2 = adam_b2
+        self.lr_decay = lr_decay
+        self.seed = seed
 
 
 if __name__ == '__main__':

toolbox/torchaudio/models/nx_clean_unet/yaml/config.yaml

@@ -22,5 +22,9 @@ discriminator_in_channel: 2
 
 compress_factor: 0.3
 
+batch_size: 4
+learning_rate: 0.0005
+adam_b1: 0.8
+adam_b2: 0.99
+lr_decay: 0.99
 seed: 1234
-batch_size: 3