From 3e14cc00d8f29b76bd0f756f2892af31e5f719ee Mon Sep 17 00:00:00 2001 From: chensiyuan Date: Thu, 16 Nov 2023 10:41:01 +0800 Subject: [PATCH] fix the self_attn module of DecoderLayer --- lib/model_zoo/seecoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/model_zoo/seecoder.py b/lib/model_zoo/seecoder.py index 7de4be0..cd0a211 100644 --- a/lib/model_zoo/seecoder.py +++ b/lib/model_zoo/seecoder.py @@ -67,7 +67,7 @@ def __init__(self, super().__init__() - self.self_attn = nn.MultiheadAttention(dim, n_heads, dropout=dropout) + self.self_attn = nn.MultiheadAttention(dim, n_heads, dropout=dropout, batch_first=True) self.dropout1 = nn.Dropout(dropout) self.norm1 = nn.LayerNorm(dim)