Repository: pytorch
Issue: [MPS] Wrong Calculation on MPS
Labels: triaged,module: correctness (silent),module: mps

### 🐛 Describe the bug

I train a single-layer preceptron, send it to the GPU (MPS), then run a loop with training, Loss starts tending to infinity (eventually becomes nan), Accuracy stays the same throughout the entire loop phase
If I run training on cpu or cuda, there is no problem, Loss goes to 0, Accuracy is more than 50%.
So I would like to point out that this is due to some internal errors of the library

Training cycle:
`for epoch in tqdm(range(100)):
    print(f"Epoch: {epoch}\n")
    train_loss = 0
    model.train()
    for batch, (X, y) in enumerate(train_dataloader):
        y_pred = model(X.to(device))
        loss = loss_fn(y_pred, y.to(device))
        train_loss += loss
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
    train_loss /= len(train_dataloader)
    test_loss, test_acc = 0, 0
    with torch.no_grad():
        for X, y in test_dataloader:
            test_pred = model(X.to(device))
            pred_classes = torch.argmax(m(test_pred), dim=1)
            true_classes = torch.argmax(y,dim=1)
            test_loss += loss_fn(test_pred, y.to(device))
            test_acc += accuracy_fn(pred_classes, true_classes.to(device))

        test_loss /= len(test_dataloader)
        test_acc /= len(test_dataloader)
    print(
        f"\nTrain loss: {train_loss:.5f} | Test loss: {test_loss:.5f}, Test acc: {test_acc:.2f}%\n"
    )`
    
 Full code to run:
[Neuron.ipynb.zip](https://github.com/pytorch/pytorch/files/14625057/Neuron.ipynb.zip)

I'm just getting trained in machine learning, but I don't think changing the device should have any effect on the results
<img width="750" alt="image" src="https://github.com/pytorch/pytorch/assets/90915334/72374625-af2a-4c2f-8cd6-6f530cd14016">
<img width="634" alt="image" src="https://github.com/pytorch/pytorch/assets/90915334/b8af3ee7-1e65-490f-85ac-f97e78de6596">



### Versions

PyTorch version: 2.2.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.4 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: version 3.27.6
Libc version: N/A

Python version: 3.11.7 | packaged by conda-forge | (main, Dec 15 2023, 08:41:36) [Clang 16.0.6 ] (64-bit runtime)
Python platform: macOS-14.4-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] flake8==6.1.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.2
[pip3] numpydoc==1.6.0
[pip3] pytorch-lightning==2.1.3
[pip3] torch==2.2.1
[pip3] torchaudio==2.2.1
[pip3] torchmetrics==1.2.1
[pip3] torchvision==0.15.2a0
[conda] numpy                     1.26.2          py311h6d074dd_0    conda-forge
[conda] numpydoc                  1.6.0              pyhd8ed1ab_0    conda-forge
[conda] pytorch                   2.2.1                  py3.11_0    pytorch
[conda] pytorch-lightning         2.1.3              pyhd8ed1ab_0    conda-forge
[conda] torchaudio                2.2.1                 py311_cpu    pytorch
[conda] torchmetrics              1.2.1              pyhd8ed1ab_0    conda-forge
[conda] torchvision               0.15.2          cpu_py311he74fb5d_0  

cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr