I have the following code for keras:
# Build Network
network = models.Sequential()
network.add(layers.Dense(392, activation='relu', input_shape=(784,)))
network.add(layers.Dense(10, activation='softmax'))
network.compile(optimizer=optimizers.RMSprop(lr=0.001),
                loss=losses.MSE,
                metrics=['acc'])

# Train
history = network.fit(train_img[:-10000], train_lbl[:-10000], epochs=10, batch_size=196,
                      validation_data=(train_img[-10000:], train_lbl[-10000:]))

Train on 50000 samples, validate on 10000 samples
Epoch 1/10
50000/50000 [==============================] - 1s 22us/step - loss: 0.0349 - acc: 0.7554 - val_loss: 0.0487 - val_acc: 0.6752
Epoch 2/10
50000/50000 [==============================] - 0s 8us/step - loss: 0.0228 - acc: 0.8425 - val_loss: 0.0317 - val_acc: 0.7911
Epoch 3/10
50000/50000 [==============================] - 0s 8us/step - loss: 0.0200 - acc: 0.8636 - val_loss: 0.0327 - val_acc: 0.7746
Epoch 4/10
50000/50000 [==============================] - 0s 8us/step - loss: 0.0182 - acc: 0.8753 - val_loss: 0.0263 - val_acc: 0.8243
Epoch 5/10
50000/50000 [==============================] - 0s 8us/step - loss: 0.0171 - acc: 0.8827 - val_loss: 0.0219 - val_acc: 0.8502
Epoch 6/10
50000/50000 [==============================] - 0s 8us/step - loss: 0.0163 - acc: 0.8881 - val_loss: 0.0175 - val_acc: 0.8805
Epoch 7/10
50000/50000 [==============================] - 0s 8us/step - loss: 0.0154 - acc: 0.8947 - val_loss: 0.0214 - val_acc: 0.8548
Epoch 8/10
50000/50000 [==============================] - 0s 8us/step - loss: 0.0149 - acc: 0.8999 - val_loss: 0.0256 - val_acc: 0.8334
Epoch 9/10
50000/50000 [==============================] - 0s 8us/step - loss: 0.0144 - acc: 0.9033 - val_loss: 0.0230 - val_acc: 0.8444
Epoch 10/10
50000/50000 [==============================] - 0s 8us/step - loss: 0.0139 - acc: 0.9070 - val_loss: 0.0182 - val_acc: 0.

I would like as close as possible equivalent in pytorch, below is my attempt:
class Sequential(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden1 = nn.Linear(784, 392)
        self.hidden2 = nn.Linear(392, 10)

    def forward(self, x):
        x = F.relu(self.hidden1(x))
        x = F.relu(self.hidden2(x))
        x = F.log_softmax(x, dim=1)
        return x

# Build Network
network = Sequential().cuda()
network

Sequential(
  (hidden1): Linear(in_features=784, out_features=392, bias=True)
  (hidden2): Linear(in_features=392, out_features=10, bias=True)
)

optimizer_fn = optim.RMSprop(network.parameters(), lr=0.001)
loss_fn = nn.MSELoss()

# Train
epochs = 10
for epoch in range(1, epochs+1):
    # forward pass
    trained = network(train_img[:-10000])
    loss = loss_fn(trained, train_lbl[:-10000])
    # validation
    val_trained = network(train_img[-10000:])
    val_loss = loss_fn(val_trained, train_lbl[-10000:])
    
    print('Epoch {epoch}/{epochs}. loss: {loss}, val_loss: {val_loss}'.format(**vars()))
    
    # backward pass
    optimizer_fn.zero_grad()
    loss.backward()
    optimizer_fn.step()

Epoch 1/10. loss: 5.864108562469482, val_loss: 5.864132404327393
Epoch 2/10. loss: 6.595648765563965, val_loss: 6.6156511306762695
Epoch 3/10. loss: 5.867150783538818, val_loss: 5.86713981628418
Epoch 4/10. loss: 5.862150192260742, val_loss: 5.862164497375488
Epoch 5/10. loss: 12.059815406799316, val_loss: 12.162103652954102
Epoch 6/10. loss: 5.862414360046387, val_loss: 5.862414836883545
Epoch 7/10. loss: 5.862414360046387, val_loss: 5.862414836883545
Epoch 8/10. loss: 5.862414360046387, val_loss: 5.862414836883545
Epoch 9/10. loss: 5.862414360046387, val_loss: 5.862414836883545
Epoch 10/10. loss: 5.862414360046387, val_loss: 5.862414836883545

I’m obviously doing something wrong. Could someone put me back on track?
Sorry if this question is too basic for these forums.