update

toolbox/torchaudio/models/nx_clean_unet/configuration_nx_clean_unet.py

@@ -8,6 +8,8 @@ class NXCleanUNetConfig(PretrainedConfig):
     https://github.com/yxlu-0102/MP-SENet/blob/main/config.json
     """
     def __init__(self,
+                 sample_rate: int = 8000,
+                 segment_size: int = 16000,
                  n_fft: int = 512,
                  win_length: int = 200,
                  hop_length: int = 80,
@@ -38,6 +40,8 @@ class NXCleanUNetConfig(PretrainedConfig):
                  **kwargs
                  ):
         super(NXCleanUNetConfig, self).__init__(**kwargs)
+        self.sample_rate = sample_rate
+        self.segment_size = segment_size
         self.n_fft = n_fft
         self.win_length = win_length
         self.hop_length = hop_length