[FIXED] SNN that works on CPU throws "Mismatched Tensor types in NNPack convolutionOutput" on MPS

Issue

As part of learning PyTorch, I created the MNIST classifier described in chapter 17 of Andrew Glassner's book Deep Learning: A Visual Approach.

I can train the network setting the device to "cpu"; however, when I set the device to "mps", nn/modules/conv.py throws the RuntimeError "Mismatched Tensor types in NNPack convolutionOutput".

Any ideas what I'm doing that might be causing this?

Here is the neural network I'm using:

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.mnist_nn = nn.Sequential(
            nn.Conv2d(1, 32, 3),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3),
            nn.ReLU(),
            nn.MaxPool2d(2,stride=2),
            nn.Dropout(p=0.25),
            nn.Flatten(),
            nn.Linear(9216, 128),
            nn.ReLU(),
            nn.Dropout(p=0.25),
            nn.Linear(128, 10),
            nn.Softmax(dim=0)
        )

    def forward(self, x):
        logits = self.mnist_nn(x)
        return logits

And the dataset I'm using is torchvision.datasets.MNIST

Thanks for your help.

Solution

I had the same problem. Make sure your data is on the same device as your model. I solved it by passing each image/label pair to the device during my training loop. That should convert it to the correct tensor-type. There are probably more intelligent ways of doing it depending on your hardware (ie. pre-loading all the data into your GPU), but this worked across devices and backends for me.

# Set the device
device = torch.device('mps')


# Within each epoch, pass your batch-data to the correct device
for img, label in training_set:
    img = img.to(device)
    label = label.to(device)

    #Continue the training-loop as usual

Answered By - Beepbo0p