Receiving 0 Acccuracy in MNIST CNN with Macbook M1 MPS

Question

I am very new to deep learning and trying to run the Pytorch official example code on MNIST dataset

code: https://github.com/pytorch/examples/tree/master/mnist

I am using Macbook Pro M1, when the device variable in Pytorch is set to mps(line 107), I got an accuracy of 0% but I can see the loss is decreasing during training

Train Epoch: 1 [56320/60000 (94%)] Loss: 0.040613
Train Epoch: 1 [56960/60000 (95%)] Loss: 0.072618
Train Epoch: 1 [57600/60000 (96%)] Loss: 0.106226
Train Epoch: 1 [58240/60000 (97%)] Loss: 0.096637
Train Epoch: 1 [58880/60000 (98%)] Loss: 0.003050
Train Epoch: 1 [59520/60000 (99%)] Loss: 0.007105

Test set: Average loss: 0.0481, Accuracy: 0/10000 (0%)

However, when I switch the device to cpu, everything works fine(but slower), got accuracy around 97%.

Does anyone know why?

Answer 1

This issue is caused by argmax(). I think it is an unsafe ops when using mps device. In this case, you can use .max(dim=1).indices to get the correct results. More details can be found at https://github.com/pytorch/pytorch/issues/92311 .