[Python] PyTorch on M4 Mac on MPS run dramatically slower compared to CPU

I'm using a M4 MacBook Pro and I'm trying to run a simple NN on MNIST data. The performance on mps is supposed to be better than that of cpu. But it is dramatically slower. Even for a simple NN like the one below, on CPU it takes around 1s, but on mps it takes ~8s. Am I missing something?

def fit(X, Y, epochs, model, optimizer):
for epoch in range(epochs):
y_pred = model.forward(X)

loss = F.binary_cross_entropy(y_pred, Y)

optimizer.zero_grad() # zero the gradients
loss.backward() # Compute new gradients
optimizer.step() # update the parameters (weights)

if (epoch % 2000 == 0):
print(f'Epoch: {epoch} | Loss: {loss.item()}')

class NeuralNet(nn.Module):
def __init__(self):
super().__init__()

self.fc1 = nn.Linear(X.shape[1], 3)
self.fc2 = nn.Linear(3, 1)

def forward(self, x):
x = F.sigmoid(self.fc1(x))
x = F.sigmoid(self.fc2(x))
return x

def predict(self, x):
output = self.forward(x)
return (output > 0.5).int()

model = NeuralNet().to(device=device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

Continue reading...