PyTorch на M4 Mac на MPS работает значительно медленнее по сравнению с процессором ⇐ Python

[Anonymous]

Я использую MacBook Pro M4 и пытаюсь запустить простую NN на данных MNIST. Производительность MPS должна быть выше, чем у процессора. Но это существенно медленнее. Даже для простой NN, такой как показанная ниже, на процессоре это занимает около 1 с, а на mps — ~ 8 с. Я что-то упустил?
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)


Подробнее здесь: https://stackoverflow.com/questions/795 ... red-to-cpu