Repository: pytorch
Issue: mps bug: failed assertion `[MPSNDArrayDescriptor sliceDimension:withSubrange:] error: subRange.start (8192) is not less than length of dimension[1] (512)'
Labels: module: crash,triaged,module: mps

### 🐛 Describe the bug

I found a bug occurred when I iterated over batches with "mps" device (it does not occur with "cpu"). The error message is similar to #96153 but I am not sure they are the same.

Here is the code to reproduce the bug.
```python
import torch 

class test_model(torch.nn.Module):
    def __init__(self, dim_in, dim_out):
        super().__init__()
        self.features = torch.nn.Sequential(
            torch.nn.Linear(dim_in, 50),
            torch.nn.LeakyReLU(),
            torch.nn.Linear(50, dim_out),
        )
    def forward(self, x):
        return self.features(x)
    

def neural_net(x, y, num_epochs = 10, batch_size = 512,):
    
    device = "mps"
    
    # make pytorch variables
    x = torch.Tensor(x).to(device)
    y = torch.Tensor(y).to(device)
    batch_size = int(batch_size)

    # initiate model
    model = test_model(x.shape[1], y.shape[1]).to(device)
    model.train()

    # train in batches
    num_samples_train = x.shape[0]
    num_batches_train = num_samples_train // batch_size

    x_train_batches = [x[i * batch_size : (i+1) * batch_size, :].to(device) for i in range(num_batches_train)]

    # train the network
    for e in range(num_epochs):
        # for each batch, calculate the gradient with respect to the loss
        for i in range(num_batches_train):
            y_pred_train = model(x_train_batches[i]) #.to(device)
            print(f'iter {e}, batch {i}')

x = torch.zeros(100000, 3) # num_samples, dim_in
y = torch.zeros(100000, 3) # num_samples, dim_out

neural_net(x, y, num_epochs = 10, batch_size = 512,)

```

The output is
```
iter 0, batch 0
/AppleInternal/Library/BuildRoots/90c9c1ae-37b6-11ee-a991-46d450270006/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSCore/Types/MPSNDArray.mm:120: failed assertion `[MPSNDArrayDescriptor sliceDimension:withSubrange:] error: subRange.start (8192) is not less than length of dimension[1] (512)'
[1]    17778 abort      python test.py
```

Note the first iteration works fine, the error happens excatly at the second iteration.


### Versions

Collecting environment information...
PyTorch version: 2.1.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.0 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.0.40.1)
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.6 | packaged by conda-forge | (main, Oct  3 2023, 10:37:07) [Clang 15.0.7 ] (64-bit runtime)
Python platform: macOS-14.0-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 M2 Max

Versions of relevant libraries:
[pip3] numpy==1.26.0
[pip3] torch==2.1.0
[conda] Could not collect

cc @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen @razarmehr