update

toolbox/torchaudio/models/clean_unet/modeling_clean_unet.py

@@ -95,17 +95,11 @@ class DecoderBlock(nn.Module):
     def forward(self, inputs: torch.Tensor):
         # inputs shape: [batch_size, channel, num_samples]
         x = self.conv(inputs)
-        print(x._version)
-
         x = self.glu(x)
-        print(x._version)
-
         x = self.convt(x)
-        print(x._version)
-        # if self.do_relu:
-        #     x = F.relu(x)
+        if self.do_relu:
+            x = F.relu(x)
             # x = self.relu(x)
-        # print(x._version)
 
         return x
 
@@ -258,7 +252,7 @@ class CleanUNet(nn.Module):
         # decoder
         for i, upsampling_block in enumerate(self.decoder):
             skip_i = skip_connections[i]
-            x += skip_i[:, :, :x.shape[-1]]
+            x = x + skip_i[:, :, :x.shape[-1]]
             x = upsampling_block(x)
 
         x = x[:, :, :L] * std