PackedSequence failure with MPS
Labels: triaged, module: mps

### 🐛 Describe the bug

Hi,

I have noticed recently that PyTorch fails in a specific way when using the PackedSequence Class with MPS. I've replicated the issue using a fresh minimal conda environment. I'm using PyTorch nightly because of the (recently closed) issues around LSTM and MPS (#96416) in PyTorch 2.0.0. The error doesn't seem to occur with PyTorch 1.13.1 (but I do get the old aten::nonzero flag with this version, so I don't know whether this version works only because the 'required' bit of code is being fed through the CPU instead of MPS).

The following code should hopefully replicate the problem I'm having. It generates some random sparsely populated time-series data (to replicate a series of padded samples), which is then packed via pack_padded_sequence, and fed into LSTM. The error occurs with larger sample sizes (like 1000), but not with smaller sample sizes (like 50).

```
import torch
from torch import nn, Tensor
from torch.nn.utils.rnn import pack_padded_sequence


# Make some sparsely populated synthetic data,
# mimicking padded time-series data
samples = 1000

input_data = torch.zeros((samples, 50, 3))
p = torch.rand((samples, 50, 3))
new_values = torch.normal(5, 3, size=(samples, 50, 3))
idx = torch.where(p < 0.05)
input_data[idx] = new_values[idx]
idx = input_data.sum(2).sum(1).nonzero().flatten()
input_data = input_data[idx]

# Set to mps
input_data = input_data.to('mps')

# Pack the data
lengths = Tensor([1 + arr.nonzero()[:, 0].max().item() - arr.nonzero()[:, 0].min().item()
                    for arr in input_data.cpu()])

packed_data = pack_padded_sequence(
            input_data,
            lengths,
            batch_first=True,
            enforce_sorted=False)

# Make our neural network
net = nn.Sequential(
    nn.LSTM(input_size=3,
            hidden_size=50)).to('mps')

output, (hn, cn) = net(packed_data)
print(output)

```

When the code is run, the following error is given, at the point of feeding the PackedSequence into the LSTM:

```
/AppleInternal/Library/BuildRoots/c651a45f-806e-11ed-a221-7ef33c48bc85/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSCore/Types/MPSNDArray.mm:82: failed assertion `[MPSNDArrayDescriptor sliceDimension:withSubrange:] error: subRange.start (8192) is not less than length of dimension[1] (995)'
```

I'm not sure I fully understand what is going wrong. The code runs fine 1) with cpu and 2) with regular padded sequence (i.e., no packing) for much larger sample sizes.

### Versions

PyTorch version: 2.1.0.dev20230324
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.2.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.0 (clang-1400.0.29.202)
CMake version: version 3.25.2
Libc version: N/A

Python version: 3.10.9 | packaged by conda-forge | (main, Feb  2 2023, 20:26:08) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.2.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Versions of relevant libraries:
[pip3] numpy==1.24.2
[pip3] torch==2.1.0.dev20230324
[conda] numpy                     1.24.2          py310h3d2048e_0    conda-forge
[conda] pytorch                   2.1.0.dev20230324        py3.10_0    pytorch-nightly

cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev