update

examples/nx_clean_unet/run.sh

@@ -12,11 +12,10 @@ 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 nx-clean-unet-aishell-20250228 \
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name nx-clean-unet-aishell-20250228 \
 --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
 
 
 sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name mpnet-nx-speech-20250224 \

toolbox/torchaudio/models/nx_clean_unet/modeling_nx_clean_unet.py

@@ -215,7 +215,9 @@ class NXCleanUNet(nn.Module):
         bottle_neck = torch.transpose(bottle_neck, dim0=-2, dim1=-1)
         # bottle_neck shape: [batch_size, time_steps, input_size]
 
+        bottle_neck = bottle_neck.unsqueeze(dim=1)
         bottle_neck = self.causal_encoder.forward(bottle_neck)
+        bottle_neck = bottle_neck.squeeze(dim=1)
         # bottle_neck shape: [batch_size, time_steps, input_size]
 
         bottle_neck = self.transformer.forward(bottle_neck)
@@ -254,7 +256,9 @@ class NXCleanUNet(nn.Module):
         bottle_neck = torch.transpose(bottle_neck, dim0=-2, dim1=-1)
         # bottle_neck shape: [batch_size, time_steps, input_size]
 
+        bottle_neck = bottle_neck.unsqueeze(dim=1)
         bottle_neck = self.causal_encoder.forward_chunk_by_chunk(bottle_neck)
+        bottle_neck = bottle_neck.squeeze(dim=1)
         # bottle_neck shape: [batch_size, time_steps, input_size]
 
         bottle_neck = self.transformer.forward_chunk_by_chunk(bottle_neck)