Update nn.Softmax for pytorch 0.4

Specify dim when initializing `nn.Softmax`
ctr4si · May 2, 2018 · be7bdd9 · be7bdd9
1 parent da0dafe
commit be7bdd9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/model/layers/decoder.py b/model/layers/decoder.py
@@ -70,7 +70,7 @@ def decode(self, out):
         # Sample next word from multinomial word distribution
         if self.sample:
             # x: [batch_size] - word index (next input)
-            x = torch.multinomial(self.softmax(out / self.temperature, dim=1), 1).view(-1)
+            x = torch.multinomial(self.softmax(out / self.temperature), 1).view(-1)
 
         # Greedy sampling
         else:
@@ -234,7 +234,7 @@ def __init__(self, vocab_size, embedding_size,
                                hidden_size,
                                dropout)
         self.out = nn.Linear(hidden_size, vocab_size)
-        self.softmax = nn.Softmax()
+        self.softmax = nn.Softmax(dim=1)
 
     def forward_step(self, x, h,
                      encoder_outputs=None,