Repository: pytorch
Issue: [JIT] jit.trace fails with custom GRUs and CUDA when the sequence is longer
Labels: oncall: jit

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
When `jit.trace` is applied to a custom GRU with CUDA enabled, the call fails when seq_len is longer than 200. 

With CUDA, the call succeeds with smaller seq_len such as 100 or 120. 
With CPU, the call succeeds even with seq_len >= 200.

## To Reproduce

Steps to reproduce the behavior:
1. Set `cuda_id = 0 (or any device)` and `seq_len = 200` in the following code

### Error message
```
torch.jit.TracingCheckError: Tracing failed sanity checks!
Encountered an exception while running the trace with test inputs.
Exception:
        default_program(28): Error: Formal parameter space overflowed (4096 bytes max) in function kernel_399
```

### Code
```python
import torch
from torch import nn


class GRU(nn.Module):
    ''' GRU
    Reference
    https://discuss.pytorch.org/t/implementation-of-multiplicative-lstm/2328/9
    https://github.com/pytorch/pytorch/blob/master/torch/nn/_functions/rnn.py#L46
    '''

    def __init__(self, input_size, hidden_size, seq_len, batch_first=True):
        super(GRU, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.batch_first = batch_first

        self.seq_len = seq_len

        self.input_weights = nn.Linear(input_size, 3 * hidden_size)
        self.hidden_weights = nn.Linear(hidden_size, 3 * hidden_size)

    def step(self, input, hidden):
        hx = hidden
        w_ih = self.input_weights(input)
        w_hh = self.hidden_weights(hx)

        i_r, i_c, i_u = w_ih.chunk(3, 1)
        h_r, h_c, h_u = w_hh.chunk(3, 1)

        updategate = torch.sigmoid(i_u+h_u)
        resetgate = torch.sigmoid(i_r+h_r)

        new_state = torch.tanh(i_c + resetgate * h_c)

        hy = (1.-updategate)*hx + updategate*new_state

        return hy

    def forward(self, input, hidden):

        if self.batch_first:
            input = input.transpose(0, 1)

        # Main loop
        output = []
        for i in range(self.seq_len):
            hidden = self.step(input[i], hidden)
            output.append(hidden)

        output = torch.cat(output, 0).view(input.size(0), *output[0].size())

        if self.batch_first:
            output = output.transpose(0, 1)

        return output, hidden


if __name__ == '__main__':
    input_size = 48
    hidden_size = 48
    batch_size = 20

    cuda_id = 0
    # cuda_id = None

    seq_len = 200
    # seq_len = 100

    gru = GRU(input_size, hidden_size, seq_len, batch_first=True)

    input = torch.rand(batch_size, seq_len, input_size)
    hidden = torch.rand(batch_size, hidden_size)

    if cuda_id is not None:
        torch.cuda.set_device(cuda_id)
        gru = gru.cuda()
        input = input.cuda()
        hidden = hidden.cuda()

    traced_gru = torch.jit.trace(gru, (input, hidden))
```
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior
`jit.trace` call succeeds.
<!-- A clear and concise description of what you expected to happen. -->

## Environment

PyTorch version: 1.0.0
Is debug build: No
CUDA used to build PyTorch: 9.0.176

OS: Ubuntu 16.04.3 LTS
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.5) 5.4.0 20160609
CMake version: version 3.5.1

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: 9.0.176
GPU models and configuration:
GPU 0: GeForce GTX 1080 Ti
GPU 1: GeForce GTX 1080 Ti

Nvidia driver version: 384.81
cuDNN version: Probably one of the following:
/usr/local/cuda-9.0/lib64/libcudnn.so
/usr/local/cuda-9.0/lib64/libcudnn.so.7
/usr/local/cuda-9.0/lib64/libcudnn.so.7.0.5
/usr/local/cuda-9.0/lib64/libcudnn_static.a

Versions of relevant libraries:
[pip] Could not collect
[conda] blas                      1.0                         mkl
[conda] mkl                       2018.0.3                      1
[conda] mkl_fft                   1.0.6            py37h7dd41cf_0
[conda] mkl_random                1.0.1            py37h4414c95_1
[conda] pytorch                   1.0.0           py3.7_cuda9.0.176_cudnn7.4.1_1    pytorch
[conda] torchvision               0.2.1                    py37_1    pytorch