Refactored again, but it still doesn’t seem right.
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 = self.hidden2(x)
        x = F.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 = 20
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/20. loss: 0.09023910015821457, val_loss: 0.09023372828960419
Epoch 2/20. loss: 0.08889125287532806, val_loss: 0.08902137726545334
Epoch 3/20. loss: 0.09533119946718216, val_loss: 0.0944860428571701
Epoch 4/20. loss: 0.11828577518463135, val_loss: 0.11910293251276016
Epoch 5/20. loss: 0.11870171129703522, val_loss: 0.11937540769577026
Epoch 6/20. loss: 0.14534510672092438, val_loss: 0.14685864746570587
Epoch 7/20. loss: 0.10830487310886383, val_loss: 0.11026407033205032
Epoch 8/20. loss: 0.12849687039852142, val_loss: 0.12968237698078156
Epoch 9/20. loss: 0.1687696874141693, val_loss: 0.17019768059253693
Epoch 10/20. loss: 0.1339481770992279, val_loss: 0.13542285561561584
Epoch 11/20. loss: 0.11552810668945312, val_loss: 0.11694050580263138
Epoch 12/20. loss: 0.13507775962352753, val_loss: 0.1372336596250534
Epoch 13/20. loss: 0.11315662413835526, val_loss: 0.11472103744745255
Epoch 14/20. loss: 0.12789055705070496, val_loss: 0.12905439734458923
Epoch 15/20. loss: 0.131318137049675, val_loss: 0.13293345272541046
Epoch 16/20. loss: 0.11097519099712372, val_loss: 0.11244189739227295
Epoch 17/20. loss: 0.10206971317529678, val_loss: 0.1033577173948288
Epoch 18/20. loss: 0.11592703312635422, val_loss: 0.11761631071567535
Epoch 19/20. loss: 0.09888488054275513, val_loss: 0.10024408251047134
Epoch 20/20. loss: 0.10537845641374588, val_loss: 0.1073707640171051

FYI, if it matters, I’m using: https://www.kaggle.com/zalando-research/fashionmnist 2