From e6ac8881f16969fefd5b151074b2e7f710ef1545 Mon Sep 17 00:00:00 2001
From: 0x45f <wangzhen45@baidu.com>
Date: Tue, 19 Jul 2022 03:43:50 +0000
Subject: [PATCH] Fix comments

---
 paddlespeech/s2t/models/u2/u2.py          | 6 ++++--
 paddlespeech/s2t/modules/encoder.py       | 6 +++---
 paddlespeech/s2t/modules/encoder_layer.py | 1 +
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/paddlespeech/s2t/models/u2/u2.py b/paddlespeech/s2t/models/u2/u2.py
index c77501848..76f698e64 100644
--- a/paddlespeech/s2t/models/u2/u2.py
+++ b/paddlespeech/s2t/models/u2/u2.py
@@ -625,10 +625,12 @@ class U2BaseModel(ASRInterface, nn.Layer):
                 (elayers, head, cache_t1, d_k * 2), where
                 `head * d_k == hidden-dim` and
                 `cache_t1 == chunk_size * num_decoding_left_chunks`.
-                `d_k * 2` for att key & value.
+                `d_k * 2` for att key & value. Default is 0-dims Tensor, 
+                it is used for dy2st.
             cnn_cache (paddle.Tensor): cache tensor for cnn_module in conformer,
                 (elayers, b=1, hidden-dim, cache_t2), where
-                `cache_t2 == cnn.lorder - 1`
+                `cache_t2 == cnn.lorder - 1`. Default is 0-dims Tensor, 
+                it is used for dy2st.
 
         Returns:
             paddle.Tensor: output of current input xs,
diff --git a/paddlespeech/s2t/modules/encoder.py b/paddlespeech/s2t/modules/encoder.py
index ad73f5e99..bff2d69bb 100644
--- a/paddlespeech/s2t/modules/encoder.py
+++ b/paddlespeech/s2t/modules/encoder.py
@@ -250,11 +250,11 @@ class BaseEncoder(nn.Layer):
         r_cnn_cache = []
         for i, layer in enumerate(self.encoders):
             # att_cache[i:i+1] = (1, head, cache_t1, d_k*2)
-            # cnn_cache[i] = (B=1, hidden-dim, cache_t2)
+            # cnn_cache[i:i+1] = (1, B=1, hidden-dim, cache_t2)
             xs, _, new_att_cache, new_cnn_cache = layer(
                 xs, att_mask, pos_emb,
-                att_cache=att_cache if elayers == 0 else att_cache[i:i+1],
-                cnn_cache=cnn_cache if paddle.shape(cnn_cache)[0] == 0 else cnn_cache[i],
+                att_cache=att_cache[i:i+1] if elayers > 0 else att_cache,
+                cnn_cache=cnn_cache[i:i+1] if paddle.shape(cnn_cache)[0] > 0 else cnn_cache,
             )
             # new_att_cache = (1, head, attention_key_size, d_k*2)
             # new_cnn_cache = (B=1, hidden-dim, cache_t2)
diff --git a/paddlespeech/s2t/modules/encoder_layer.py b/paddlespeech/s2t/modules/encoder_layer.py
index d91e3f6ef..9e46cc540 100644
--- a/paddlespeech/s2t/modules/encoder_layer.py
+++ b/paddlespeech/s2t/modules/encoder_layer.py
@@ -250,6 +250,7 @@ class ConformerEncoderLayer(nn.Layer):
         # convolution module
         # Fake new cnn cache here, and then change it in conv_module
         new_cnn_cache = paddle.zeros([0,0,0], dtype=x.dtype)
+        cnn_cache = paddle.squeeze(cnn_cache, axis=0)
         if self.conv_module is not None:
             residual = x
             if self.normalize_before: