Support calling pack_padedd_sequence with a Variable lengths (#5113)

anderspapitto · ezyang · commit ef1459020958 · 2018-02-07T17:11:33.000-05:00
This was accidentally lost while addressing review comments on #4695 pack_padded_sequence may be called either with a list or with a Variable. If called with a list we convert to Variable internally. I added to test_nn to test the new codepath. The bug was also caught by the onnx-fb-universe tests (which rely on passing in Variable).
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -3053,7 +3053,9 @@ def compare_cpu_gpu(outputs_cpu, outputs_gpu):
                     self.assertEqual(cpu_weight.grad.data, gpu_weight.grad.data, prec=5e-5)
 
         for module in (nn.RNN, nn.LSTM, nn.GRU):
-            for bias, bidirectional, batch_first, contig, variable_len in product((True, False), repeat=5):
+            for bias, bidirectional, batch_first, contig, variable_len, lens_as_variable \
+                    in product((True, False), repeat=6):
+
                 num_directions = 2 if bidirectional else 1
                 if batch_first:
                     input_val = torch.randn(batch, seq_length, input_size)
@@ -3073,6 +3075,8 @@ def compare_cpu_gpu(outputs_cpu, outputs_gpu):
 
                 if variable_len:
                     lengths = [7, 5, 5, 2, 1, 1]
+                    if lens_as_variable:
+                        lengths = Variable(torch.LongTensor(lengths))
                     input_val = Variable(input_val)
                     grad_output = Variable(grad_output)
                     input_val = rnn_utils.pack_padded_sequence(input_val, lengths, batch_first=batch_first)
diff --git a/torch/nn/_functions/packing.py b/torch/nn/_functions/packing.py
@@ -15,8 +15,8 @@ def forward(ctx, input, lengths, batch_first):
         steps = []
         batch_sizes = []
 
-        # lengths is a Tensor, so we must convert to list before reversed()
-        lengths_iter = reversed(list(lengths))
+        # lengths is a Tensor, so we must convert to [int] before reversed()
+        lengths_iter = reversed(lengths.tolist())
 
         batch_size = input.size(1)
 
diff --git a/torch/nn/utils/rnn.py b/torch/nn/utils/rnn.py
@@ -91,6 +91,9 @@ def is_cuda(self):
 
 
 def _pack_padded_sequence(input, lengths, batch_first=False):
+    if isinstance(lengths, list):
+        lengths = Variable(torch.LongTensor(lengths))
+
     data, batch_sizes = PackPadded.apply(input, lengths, batch_first)
 
     return PackedSequence(data, batch_sizes)
