update

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

@@ -70,13 +70,12 @@ class TransformerEncoderLayer(nn.Module):
             self,
             x: torch.Tensor,
             mask: torch.Tensor,
-            attention_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+            attention_cache: torch.Tensor = None,
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """
 
         :param x: torch.Tensor. shape=(batch_size, time, input_dim).
         :param mask: torch.Tensor. mask tensor for the input. shape=(batch_size, time，time).
-        :param position_embedding: torch.Tensor.
         :param attention_cache: torch.Tensor. cache tensor of the KEY & VALUE
                 shape=(batch_size=1, head, cache_t1, d_k * 2), head * d_k == input_dim.
         :return:
@@ -177,8 +176,8 @@ class TransformerEncoder(nn.Module):
     def forward_chunk(self,
                       xs: torch.Tensor,
                       offset: int,
-                      attention_mask: torch.Tensor = torch.zeros(0, 0, 0),
-                      attention_cache: torch.Tensor = torch.zeros(0, 0, 0, 0),
+                      attention_mask: torch.Tensor = None,
+                      attention_cache: torch.Tensor = None,
                       ) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         Forward just one chunk.