From 407a44a5441c86f4f0cf2283aca75a34a348249d Mon Sep 17 00:00:00 2001
From: duzx16 <zx-du20@mails.tsinghua.edu.cn>
Date: Sun, 12 Dec 2021 22:14:48 +0800
Subject: [PATCH] Fix return values for SelfAttention
---
SwissArmyTransformer/mpu/transformer.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/SwissArmyTransformer/mpu/transformer.py b/SwissArmyTransformer/mpu/transformer.py
index d9d3181..3b4ae73 100755
--- a/SwissArmyTransformer/mpu/transformer.py
+++ b/SwissArmyTransformer/mpu/transformer.py
@@ -158,7 +158,7 @@ class SelfAttention(torch.nn.Module):
if self.training:
output = self.output_dropout(output)
- return output, None
+ return output
class CrossAttention(torch.nn.Module):
--
GitLab