Repository: pytorch
Issue: Network training silently fails with MPS backend
Labels: needs reproduction,triaged,module: correctness (silent),module: mps

### 🐛 Describe the bug

I'm working with a small conv1D network. The network trains as expected when running the script with a CPU (or CUDA) backend. But if the backend is switched to MPS, the network training seems to fail silently (no errors are raised, but the network doesn't learn anything as it did with CPU or CUDA backend). The difference between the network in the script trained using MPS vs. CPU are substantial, and even more so if the only ReLu layer in the network is removed (see code and output snippets below).

I tried to whittle down the code to isolate what is causing the issue, but wasn't really able to do that too well. The code (and the data required to train it) are on [this](https://gist.github.com/UjasShah/c7785e79ae1049f56ce5c39643498a00) public gist (you might have to scroll down quite a bit for the code). The code really just is 1) manipulating the data a bit to make it as the network expects it,  2) the network itself, 3) a masked cross entropy loss function and 4) some optional code to produce outputs from a trained network.

The data can also be gotten from [here](https://github.com/karpathy/makemore/blob/master/names.txt) 

I've pasted the network training below to show the difference between MPS and CPU. Shown below are differences I've seen consistently over many runs.
CPU:
```
Using cpu device
step: 0, train loss: 3.22050, val loss: 3.23820
step: 1000, train loss: 2.15283, val loss: 2.39201
```
MPS:
```
Using mps device
step: 0, train loss: 3.23463, val loss: 3.25037
step: 1000, train loss: 2.80399, val loss: 2.92336
```

Another example of networks training differently is when I remove the only ReLu layer in the network
CPU:
```
Using cpu device
step: 0, train loss: 3.10652, val loss: 3.13682
step: 1000, train loss: 2.25550, val loss: 2.44400
```
MPS:
```
Using mps device
step: 0, train loss: 3.24397, val loss: 3.26047
step: 1000, train loss: 1952947712.00000, val loss: 2297497600.00000
```

Sorry for the long code! As I said, I tried to get to the exact operation causing this bug but wasn't able to. This is also my first GitHub issue so any pointers on making this issue better or triaging the bug would be welcome!

Also wanted to link [another](https://github.com/pytorch/pytorch/issues/137001) similar issue I saw (might not be a lot of help though).

### Versions

```
Collecting environment information...
PyTorch version: 2.4.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 15.0.1 (arm64)
GCC version: Could not collect
Clang version: 16.0.0 (clang-1600.0.26.3)
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.6 (default, Aug  9 2024, 14:24:13)  [Clang 16.0.0 (clang-1600.0.26.3)] (64-bit runtime)
Python platform: macOS-15.0.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 M3

Versions of relevant libraries:
[pip3] numpy==2.0.2
[pip3] torch==2.4.1
[pip3] torchtext==0.18.0
[pip3] torchvision==0.19.1
[conda] Could not collect
```

cc @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen