From 75a704e101d0aceb973940f0fa23485fb4909f5e Mon Sep 17 00:00:00 2001 From: huangxinjing Date: Wed, 20 Oct 2021 10:05:26 +0800 Subject: [PATCH] Fix transformer dropout error --- mindspore/parallel/nn/transformer.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mindspore/parallel/nn/transformer.py b/mindspore/parallel/nn/transformer.py index 2f6e1e33e2b..457e2cfa218 100644 --- a/mindspore/parallel/nn/transformer.py +++ b/mindspore/parallel/nn/transformer.py @@ -370,6 +370,8 @@ class FeedForward(Cell): self.projection.bias.parallel_optimizer = False self.dropout = nn.Dropout(1 - dropout_rate) self.dropout.dropout.shard(((dp, 1),)) + self.dropout_3d = nn.Dropout(1 - dropout_rate) + self.dropout_3d.dropout.shard(((dp, 1, 1),)) self.cast = P.Cast() def construct(self, x): @@ -380,7 +382,10 @@ class FeedForward(Cell): hidden = self.mapping(x) output = self.projection(hidden) # returned shape is [bs, seq_length, ffn_hidden_size] or [bs * seq_length, ffn_hidden_size] - output = self.dropout(output) + if len(F.shape(output)) == 3: + output = self.dropout_3d(output) + else: + output = self.dropout(output) return output