Fix softmaxes

linguisticsai · Apr 11, 2018 · efcfc1c · efcfc1c
1 parent 586e147
commit efcfc1c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/onmt/ModelConstructor.py b/onmt/ModelConstructor.py
@@ -191,7 +191,7 @@ def make_base_model(model_opt, fields, gpu, checkpoint=None):
     if not model_opt.copy_attn:
         generator = nn.Sequential(
             nn.Linear(model_opt.rnn_size, len(fields["tgt"].vocab)),
-            nn.LogSoftmax())
+            nn.LogSoftmax(dim=-1))
         if model_opt.share_decoder_embeddings:
             generator[0].weight = decoder.embeddings.word_lut.weight
     else:

diff --git a/onmt/modules/GlobalAttention.py b/onmt/modules/GlobalAttention.py
@@ -76,7 +76,7 @@ def __init__(self, dim, coverage=False, attn_type="dot"):
         out_bias = self.attn_type == "mlp"
         self.linear_out = nn.Linear(dim*2, dim, bias=out_bias)
 
-        self.sm = nn.Softmax()
+        self.sm = nn.Softmax(dim=-1)
         self.tanh = nn.Tanh()
 
         if coverage: