From cf93d1f12915fc5c2c7cdabc07d61c66d9d9224f Mon Sep 17 00:00:00 2001 From: Daniel Povey Date: Tue, 16 May 2023 17:29:34 +0800 Subject: [PATCH] Bug fix regarding chunk-size reshaping --- egs/libriheavy/LM/zipformer1/subformer.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/egs/libriheavy/LM/zipformer1/subformer.py b/egs/libriheavy/LM/zipformer1/subformer.py index 5b02921ea..e1aca3fdb 100644 --- a/egs/libriheavy/LM/zipformer1/subformer.py +++ b/egs/libriheavy/LM/zipformer1/subformer.py @@ -682,13 +682,12 @@ class SubformerEncoder(nn.Module): def _to_chunk_size(self, src: Tensor, chunk_size: int) -> Tensor: """ - Reshape embeddings 'src' to have a different chunk size (in frames) + Reshape embeddings 'src' to have a different chunk size (in frames) by + changing the batch size. """ (seq_len, batch_size, num_channels) = src.shape - num_chunks = src.shape[0] // chunk_size - src = src.reshape(num_chunks, chunk_size, batch_size, num_channels) - return src.contiguous().reshape(chunk_size, num_chunks * batch_size, - num_channels) + src = src.transpose(0, 1).contiguous().reshape(-1, chunk_size, num_channels) + return src.transpose(0, 1).contiguous() def _attn_offset_to_chunk_size(self, attn_offset: Tensor, batch_size: int, chunk_size: int) -> Tensor: