update

examples/nx_clean_unet/run.sh

@@ -12,7 +12,7 @@ sh run.sh --stage 3 --stop_stage 3 --system_version centos --file_folder_name fi
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
 
-sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name mpnet-aishell-20250224 \
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name mpnet-aishell-20250224 \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train" \
 --max_epochs 100 --max_count 10000

examples/nx_clean_unet/yaml/config.yaml

@@ -23,3 +23,4 @@ discriminator_in_channel: 2
 compress_factor: 0.3
 
 seed: 1234
+batch_size: 3

toolbox/torchaudio/models/nx_clean_unet/configuration_nx_clean_unet.py

@@ -29,6 +29,7 @@ class NXCleanUNetConfig(PretrainedConfig):
                  compress_factor: float = 0.3,
 
                  seed: int = 1234,
+                 batch_size: int = 3,
                  **kwargs
                  ):
         super(NXCleanUNetConfig, self).__init__(**kwargs)
@@ -53,6 +54,7 @@ class NXCleanUNetConfig(PretrainedConfig):
         self.compress_factor = compress_factor
 
         self.seed = seed
+        self.batch_size = batch_size
 
 
 if __name__ == '__main__':

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

@@ -23,3 +23,4 @@ discriminator_in_channel: 2
 compress_factor: 0.3
 
 seed: 1234
+batch_size: 3