update

toolbox/torchaudio/models/nx_clean_unet/modeling_nx_clean_unet.py

@@ -133,7 +133,6 @@ class UpSampling(nn.Module):
             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]
         for up_sampling_block in self.up_sampling_block_list:

toolbox/torchaudio/models/nx_clean_unet/transformer/transformer.py

@@ -467,6 +467,7 @@ class TransformerEncoder(nn.Module):
             chunk_size=self.chunk_size,
             num_left_chunks=self.num_left_chunks
         )
+        chunk_masks = chunk_masks.to(xs.device)
         # chunk_masks shape: [1, time_steps, time_steps]
         chunk_masks = torch.broadcast_to(chunk_masks, size=(batch_size, time_steps, time_steps))
         # chunk_masks shape: [batch_size, time_steps, time_steps]