Repository: pytorch/pytorch

Issue:
Improve error message for RuntimeError: inconsistent tensor size in RNN
I run RNN sucessfully with the code.
```
rnn = nn.RNN(features, hidden_size, num_layers)
o, hidden = rnn(x, hidden)
```
and I run embedding example. Then, got RuntimeError.
```
RuntimeError: inconsistent tensor size, expected r_ [5 x 25], t [5 x 25] and src [2 x 25] to 
have thesame number of elements, but got 125, 125 and 50 elements respectively at 
/opt/conda/conda-bld/pytorch_1503970438496/work/torch/lib/TH/generic/THTensorMath.c:887
```
I set `batch_first=True` to use `Embedding`
```
x.size is [5, 8, 15]
hidden_size is [5, 2, 25]
```
nothing seems wrong.
And the error info seems not the dim problem?
The error is so strange. 



Discussion:
What exactly do you mean by "embedding example"? Could you provide a code snippet that causes the RuntimeError?
@zou3519 
To better reproduce the error, I use the official code example of [torch.nn.RNN](http://pytorch.org/docs/master/nn.html?highlight=rnn#recurrent-layers)
The origin runs well:
```
rnn = nn.RNN(10, 20, 2, batch_first=False)
input = Variable(torch.randn(5, 3, 10))
h0 = Variable(torch.randn(2, 3, 20))
output, hn = rnn(input, h0)
output.size(), hn.size()
```
Due to use `embedding`, I have to make batch first.
And then , it raises error.
```
rnn = nn.RNN(10, 20, 2, batch_first=True)
input = Variable(torch.randn(5, 3, 10).transpose(0, 1))
h0 = Variable(torch.randn(2, 3, 20).transpose(0, 1))
output, hn = rnn(input, h0)
output.size(), hn.size()
```

```
RuntimeError: inconsistent tensor size, expected r_ [3 x 20], t [3 x 20] and src [2 x 20] to
 have the same number of elements, but got 60, 60 and 40 elements respectively at 
/opt/conda/conda-bld/pytorch_1503970438496/work/torch/lib/TH/generic/THTensorMath.c:887
```
Thank you for the example!

The [docs](http://pytorch.org/docs/master/nn.html?highlight=rnn#torch.nn.RNN) say that setting `batch_first=True` means that the first dim of input is the batch, ie, input has size `(batch, seq_len, num_features)`. `h0` however is unchanged, so the following should work:

```
rnn = nn.RNN(10, 20, 2, batch_first=True)
input = Variable(torch.randn(5, 3, 10).transpose(0, 1))
h0 = Variable(torch.randn(2, 3, 20))
output, hn = rnn(input, h0)
output.size(), hn.size()
```

I do agree the error message could be better, though

Resolving patch:
diff --git a/torch/backends/cudnn/rnn.py b/torch/backends/cudnn/rnn.py
--- a/torch/backends/cudnn/rnn.py
+++ b/torch/backends/cudnn/rnn.py
@@ -203,13 +203,6 @@ def forward(fn, input, hx, weight, output, hy):
         if fn.batch_first and not is_input_packed:
             input = input.transpose(0, 1)
 
-        if (not is_input_packed and input.dim() != 3) or (is_input_packed and input.dim() != 2):
-            raise RuntimeError(
-                'input must have 3 dimensions, got {}'.format(input.dim()))
-        if fn.input_size != input.size(-1):
-            raise RuntimeError('input.size(-1) must be equal to input_size. Expected {}, got {}'.format(
-                fn.input_size, input.size(-1)
-            ))
         if fn.dropout != 0 and cudnn.version() < 5103:
             raise RuntimeError('dropout supported only in cudnn v5.1 and above')
 
@@ -261,9 +254,6 @@ def forward(fn, input, hx, weight, output, hy):
             fn.w_desc = init_weight_descriptor(fn, fn.weight_buf)
             w = fn.weight_buf
 
-        if tuple(hx.size()) != hidden_size:
-            raise RuntimeError('Expected hidden size {}, got {}'.format(
-                hidden_size, tuple(hx.size())))
         if cx is not None and tuple(cx.size()) != hidden_size:
             raise RuntimeError('Expected cell size {}, got {}'.format(
                 hidden_size, tuple(cx.size())))
diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py
--- a/torch/nn/modules/rnn.py
+++ b/torch/nn/modules/rnn.py
@@ -128,6 +128,39 @@ def reset_parameters(self):
         for weight in self.parameters():
             weight.data.uniform_(-stdv, stdv)
 
+    def check_forward_args(self, input, hidden, batch_sizes):
+        is_input_packed = batch_sizes is not None
+        expected_input_dim = 2 if is_input_packed else 3
+        if input.dim() != expected_input_dim:
+            raise RuntimeError(
+                'input must have {} dimensions, got {}'.format(
+                    expected_input_dim, input.dim()))
+        if self.input_size != input.size(-1):
+            raise RuntimeError(
+                'input.size(-1) must be equal to input_size. Expected {}, got {}'.format(
+                    fn.input_size, input.size(-1)))
+
+        if is_input_packed:
+            mini_batch = batch_sizes[0]
+        else:
+            mini_batch = input.size(0) if self.batch_first else input.size(1)
+
+        num_directions = 2 if self.bidirectional else 1
+        expected_hidden_size = (self.num_layers * num_directions,
+                                mini_batch, self.hidden_size)
+
+        def check_hidden_size(hx, expected_hidden_size, msg='Expected hidden size {}, got {}'):
+            if tuple(hx.size()) != expected_hidden_size:
+                raise RuntimeError(msg.format(expected_hidden_size, tuple(hx.size())))
+
+        if self.mode == 'LSTM':
+            check_hidden_size(hidden[0], expected_hidden_size,
+                              'Expected hidden[0] size {}, got {}')
+            check_hidden_size(hidden[1], expected_hidden_size,
+                              'Expected hidden[1] size {}, got {}')
+        else:
+            check_hidden_size(hidden, expected_hidden_size)
+
     def forward(self, input, hx=None):
         is_packed = isinstance(input, PackedSequence)
         if is_packed:
@@ -153,6 +186,8 @@ def forward(self, input, hx=None):
             flat_weight = first_data.new().set_(first_data.storage(), 0, torch.Size([self._param_buf_size]))
         else:
             flat_weight = None
+
+        self.check_forward_args(input, hx, batch_sizes)
         func = self._backend.RNN(
             self.mode,
             self.input_size,


Test patch:
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -2423,6 +2423,69 @@ def check_rnn_grads(rnn1, rnn2):
         finally:
             torch.backends.cudnn.enabled = prev
 
+    def test_rnn_args_check(self):
+        input_size = 3
+        hidden_size = 5
+        num_layers = 2
+        batch_size = 4
+        seq_len = 6
+        num_directions = 1
+
+        def test(input_shape, hidden_shape, mode):
+            for input, hidden in get_inputs(input_shape, hidden_shape, mode):
+                model = getattr(nn, mode)(input_size, hidden_size, num_layers)
+                self.assertRaises(RuntimeError, lambda: model(input, hidden))
+
+        correct_input_shape = (seq_len, batch_size, input_size)
+        correct_hidden_shape = (num_layers * num_directions, batch_size, hidden_size)
+
+        def update_tuple(tup, dim, delta):
+            new_tup = list(tup)
+            new_tup[dim] = delta
+            return tuple(new_tup)
+
+        def get_inputs(input_shape, hidden_shape, mode):
+            '''returns list( tuple(input, hidden) )
+            where input, hidden are inputs to a model'''
+            input = Variable(torch.randn(input_shape))
+            hidden = Variable(torch.randn(hidden_shape))
+            if mode is not 'LSTM':
+                return [(input, hidden)]
+            if hidden_shape == correct_hidden_shape:
+                return [(input, (hidden, hidden))]
+            good_hidden = Variable(torch.randn(correct_hidden_shape))
+            return [
+                (input, (hidden, good_hidden)),
+                (input, (good_hidden, hidden)),
+            ]
+
+        rnn_modes = ['RNN', 'GRU', 'LSTM']
+        for mode in rnn_modes:
+            # Incorrect input batch size
+            input_shape = update_tuple(correct_input_shape, 1, -1)
+            hidden_shape = correct_hidden_shape
+            test(input_shape, hidden_shape, mode)
+
+            # Incorrect hidden batch size
+            input_shape = correct_input_shape
+            hidden_shape = update_tuple(correct_hidden_shape, 1, -1)
+            test(input_shape, hidden_shape, mode)
+
+            # Incorrect input size
+            input_shape = update_tuple(correct_input_shape, 2, -1)
+            hidden_shape = correct_hidden_shape
+            test(input_shape, hidden_shape, mode)
+
+            # Incorrect hidden size
+            input_shape = correct_input_shape
+            hidden_shape = update_tuple(correct_hidden_shape, 2, -1)
+            test(input_shape, hidden_shape, mode)
+
+            # Incorrect hidden[0]
+            input_shape = correct_input_shape
+            hidden_shape = update_tuple(correct_hidden_shape, 0, -1)
+            test(input_shape, hidden_shape, mode)
+
     def test_rnn_initial_hidden_state(self):
         rnn_modes = ['RNN', 'GRU', 'LSTM']
         for mode in rnn_modes: