Repository: pytorch
Issue: Libtorch segfault in packed GRU evaluation with cuda batch_sizes
Labels: module: cpp,module: nn,module: rnn,triaged

## 🐛 Bug
Calling [`torch::gru`](https://pytorch.org/cppdocs/api/function_namespaceat_1ade51e3fd651aba8f03ac23f95af11b62.html) for packed inputs when the `batch_sizes` tensor is on CUDA results in a segfault. I believe this happens with LSTMs too but I haven't tested that as carefully.

The requirement that `batch_sizes` should be on the CPU is documented in the [`PackedSequence` class](https://pytorch.org/cppdocs/api/classtorch_1_1nn_1_1utils_1_1rnn_1_1_packed_sequence.html#classtorch_1_1nn_1_1utils_1_1rnn_1_1_packed_sequence) and [`pack_padded_sequence`](https://pytorch.org/cppdocs/api/function_namespacetorch_1_1nn_1_1utils_1_1rnn_1a6c14a90e57eb631f51f06e52a600f7f7.html) raises an exception if `batch_sizes` is not on the CPU, but when calling `torch::gru` directly then there is no exception and a segfault occurs instead.
It would be nicer if `torch::gru` raised an exception as well instead of segfaulting.

## To Reproduce
```cpp
#include <torch/torch.h>
#include <iostream>
#include <tuple>

int main() {
  auto options = torch::dtype(torch::kFloat).device(torch::kCUDA, 0);
  torch::Tensor input = torch::zeros({10, 1}, options);
  torch::Tensor batch_sizes =
      torch::ones({10}, torch::dtype(torch::kLong).device(torch::kCUDA, 0));

  /*
  // Raises exception with CUDA batch_sizes
  torch::Tensor padded_input = torch::zeros({1, 10, 1});
  auto packed =
      torch::nn::utils::rnn::pack_padded_sequence(padded_input, batch_sizes);
  */

  torch::Tensor hx = torch::zeros({1, 1, 8}, options);
  torch::Tensor w_ih = torch::zeros({24, 1}, options);
  torch::Tensor w_hh = torch::zeros({24, 8}, options);
  std::vector<torch::Tensor> params_vec;
  params_vec.push_back(w_ih);
  params_vec.push_back(w_hh);
  torch::TensorList params = torch::TensorList(params_vec);
  // Segfault with CUDA batch_sizes
  auto result =
      torch::gru(input, batch_sizes, hx, params, false, 1, 0.0, false, false);
  std::cout << std::get<0>(result) << std::endl;
}

```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

`torch::gru` raises an exception if `batch_sizes` is not on the CPU instead of segfaulting.

## Environment
<details>

```
PyTorch version: 1.8.1
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (GCC) 11.1.0
Clang version: 11.1.0
CMake version: version 3.20.3
Libc version: glibc-2.33

Python version: 3.9 (64-bit runtime)
Python platform: Linux-5.12.6-arch1-1-x86_64-with-glibc2.33
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 960
Nvidia driver version: 465.31
cuDNN version: Probably one of the following:
/usr/lib/libcudnn.so.8.2.0
/usr/lib/libcudnn_adv_infer.so.8.2.0
/usr/lib/libcudnn_adv_train.so.8.2.0
/usr/lib/libcudnn_cnn_infer.so.8.2.0
/usr/lib/libcudnn_cnn_train.so.8.2.0
/usr/lib/libcudnn_ops_infer.so.8.2.0
/usr/lib/libcudnn_ops_train.so.8.2.0
HIP runtime version: N/A
MIOpen runtime version: N/A
```
</details>

cc @yf225 @glaringlee @albanD @mruberry @jbschlosser @zou3519