Fix on torch.argsort only supporting for >PyTorch1.0.0 (songyouwei#24)

* Fix version incompatibility of PyTorch * Update dynamic_rnn.py
ninadakolekar · Jan 26, 2019 · b80b895 · b80b895
1 parent 7dd9f35
commit b80b895
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -8,7 +8,7 @@
 
 ## Requirement
 
-* PyTorch 0.4.0
+* PyTorch >= 0.4.0
 * NumPy 1.13.3
 * tensorboardX 1.2
 * Python 3.6

diff --git a/layers/dynamic_rnn.py b/layers/dynamic_rnn.py
@@ -57,10 +57,10 @@ def forward(self, x, x_len):
         :return:
         """
         """sort"""
-        x_sort_idx = torch.argsort(-x_len)
-        x_unsort_idx = torch.argsort(x_sort_idx).long()
+        x_sort_idx = torch.sort(-x_len)[1].long()
+        x_unsort_idx = torch.sort(x_sort_idx)[1].long()
         x_len = x_len[x_sort_idx]
-        x = x[x_sort_idx.long()]
+        x = x[x_sort_idx]
         """pack"""
         x_emb_p = torch.nn.utils.rnn.pack_padded_sequence(x, x_len, batch_first=self.batch_first)