update

toolbox/torchaudio/models/nx_clean_unet/modeling_nx_clean_unet.py

@@ -50,8 +50,7 @@ class DownSampling(nn.Module):
         super(DownSampling, self).__init__()
         self.num_layers = num_layers
 
-        self.down_sampling_block_list = list()
-
+        down_sampling_block_list = list()
         for idx in range(self.num_layers):
             down_sampling_block = DownSamplingBlock(
                 in_channels=in_channels,
@@ -62,6 +61,8 @@ class DownSampling(nn.Module):
             self.down_sampling_block_list.append(down_sampling_block)
             in_channels = hidden_channels
 
+        self.down_sampling_block_list = nn.ModuleList(modules=down_sampling_block_list)
+
     def forward(self, x: torch.Tensor):
         # x shape: [batch_size, channels, num_samples]
         for down_sampling_block in self.down_sampling_block_list:
@@ -111,8 +112,7 @@ class UpSampling(nn.Module):
         super(UpSampling, self).__init__()
         self.num_layers = num_layers
 
-        self.up_sampling_block_list = list()
-
+        up_sampling_block_list = list()
         for idx in range(self.num_layers-1):
             up_sampling_block = UpSamplingBlock(
                 out_channels=hidden_channels,
@@ -131,6 +131,8 @@ class UpSampling(nn.Module):
                 do_relu=False,
             )
             self.up_sampling_block_list.append(up_sampling_block)
+        self.up_sampling_block_list = nn.ModuleList(modules=up_sampling_block_list)
+
 
     def forward(self, x: torch.Tensor):
         # x shape: [batch_size, channels, num_samples]