Repository: pytorch
Issue: Batch size is hardcoded using torch.jit.trace with LSTMCell
Labels: oncall: jit

## 🐛 Bug

Batch size is hardcoded when tracing a model using custom for loop with `nn.LSTMCell`. This makes it not possible to run model inference with different batch sizes. 

## To Reproduce

Steps to reproduce the behavior:

A toy model to reproduce the issue
```
import torch
import torch.nn as nn
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class TestLSTMCell(nn.Module):
    def __init__(self, feature_size, hidden_size):
        super(TestLSTMCell, self).__init__()
        self.hidden_size = hidden_size
        self.rnn = nn.LSTMCell(feature_size, hidden_size)
        
    def forward(self, input, hidden):
        num_steps = input.shape[0]
        batch_size = input.shape[1]
        output = torch.FloatTensor(num_steps, batch_size, self.hidden_size).fill_(0).to(device)
        for i in range(num_steps):
            hidden = self.rnn(input[i], hidden)
            output[i,:,:] = hidden[0]
        return output

feature_size = 20
hidden_size = 10
max_length = 10

model = TestLSTMCell(feature_size, hidden_size).to(device)

batch_size = 5
input0 = torch.rand(max_length, batch_size, feature_size, dtype=torch.float32).to(device)
hidden0 = (torch.FloatTensor(batch_size, hidden_size).fill_(0).to(device), torch.FloatTensor(batch_size, hidden_size).fill_(0).to(device))
traced_script_module = torch.jit.trace(model, (input0, hidden0))

batch_size = 10
input1 = torch.rand(max_length, batch_size, feature_size, dtype=torch.float32).to(device)
hidden1 = (torch.FloatTensor(batch_size, hidden_size).fill_(0).to(device), torch.FloatTensor(batch_size, hidden_size).fill_(0).to(device))
output =traced_script_module(input1, hidden1)
```
This gives the error
```
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
debug.py(182): forward
/opt/conda/envs/etr/lib/python3.7/site-packages/torch/nn/modules/module.py(860): _slow_forward
/opt/conda/envs/etr/lib/python3.7/site-packages/torch/nn/modules/module.py(887): _call_impl
/opt/conda/envs/etr/lib/python3.7/site-packages/torch/jit/_trace.py(940): trace_module
/opt/conda/envs/etr/lib/python3.7/site-packages/torch/jit/_trace.py(742): trace
debug.py(194): <module>
RuntimeError: The size of tensor a (5) must match the size of tensor b (10) at non-singleton dimension 0
```
## Expected behavior
I expect the traced model can run inference with a different batch size and return the correct results.

## Environment

PyTorch version: 1.8.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.10.2
Libc version: glibc-2.10

Python version: 3.7 (64-bit runtime)
Python platform: Linux-4.15.0-43.generic.x86_64-x86_64-with-debian-buster-sid
Is CUDA available: False
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: Tesla P100-PCIE-16GB
Nvidia driver version: 418.87.00
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.18.5
[pip3] numpydoc==1.1.0
[pip3] torch==1.8.1
[pip3] torchvision==0.2.1
[conda] blas                      1.0                         mkl  
[conda] mkl                       2021.2.0           h06a4308_296  
[conda] mkl-service               2.3.0            py37h27cfd23_1  
[conda] mkl_fft                   1.3.0            py37h42c9631_2  
[conda] mkl_random                1.2.1            py37ha9443f7_2  
[conda] numpy                     1.18.5                   pypi_0    pypi
[conda] numpydoc                  1.1.0              pyhd3eb1b0_1  
[conda] torch                     1.8.1                    pypi_0    pypi
[conda] torchvision               0.2.1                    pypi_0    pypi

## Additional context
We are trying to use a model from a research paper who developed a custom attention mechanism. The attention mechanism is implemented with a custom for loop and `nn.LSTMCell`. We got above error when tracing the whole model. After some investigations we found out that the issue comes from tracing `nn.LSTMCell` specifically. We realized using `nn.LSTM` may not have such behaviors (we tried a toy model for `nn.LSTM` and it worked fine). However as the original model comes from a 3rd-party paper and the attention mechanism is quite complicated, rewriting the network using `nn.LSTM` is not really feasible. 

cc @gmagogsfm