From 2a66366be4c2715870e4859fd5a5db6e8a9dc00a Mon Sep 17 00:00:00 2001
From: chenmengzheAAA <123789350+chenmengzheAAA@users.noreply.github.com>
Date: 星期四, 14 九月 2023 19:00:17 +0800
Subject: [PATCH] Merge pull request #956 from alibaba-damo-academy/chenmengzheAAA-patch-4
---
funasr/export/models/modules/multihead_att.py | 12 ++++++------
1 files changed, 6 insertions(+), 6 deletions(-)
diff --git a/funasr/export/models/modules/multihead_att.py b/funasr/export/models/modules/multihead_att.py
index 1983db8..4885c4e 100644
--- a/funasr/export/models/modules/multihead_att.py
+++ b/funasr/export/models/modules/multihead_att.py
@@ -64,19 +64,19 @@
return self.linear_out(context_layer) # (batch, time1, d_model)
-def preprocess_for_attn(x, mask, cache, pad_fn):
+def preprocess_for_attn(x, mask, cache, pad_fn, kernel_size):
x = x * mask
x = x.transpose(1, 2)
if cache is None:
x = pad_fn(x)
else:
- x = torch.cat((cache[:, :, 1:], x), dim=2)
- cache = x
+ x = torch.cat((cache, x), dim=2)
+ cache = x[:, :, -(kernel_size-1):]
return x, cache
-torch_version = float(".".join(torch.__version__.split(".")[:2]))
-if torch_version >= 1.8:
+torch_version = tuple([int(i) for i in torch.__version__.split(".")[:2]])
+if torch_version >= (1, 8):
import torch.fx
torch.fx.wrap('preprocess_for_attn')
@@ -90,7 +90,7 @@
self.attn = None
def forward(self, inputs, mask, cache=None):
- x, cache = preprocess_for_attn(inputs, mask, cache, self.pad_fn)
+ x, cache = preprocess_for_attn(inputs, mask, cache, self.pad_fn, self.kernel_size)
x = self.fsmn_block(x)
x = x.transpose(1, 2)
--
Gitblit v1.9.1