here is the whole network training code:
train_data = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=None)
test_data = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=None)

X_train, y_train = train_data.data, train_data.targets
X_test, y_test = test_data.data, test_data.targets
X_train = X_train.float()
X_test = X_test.float()
X_train = X_train.reshape(-1, 28 * 28)
X_test = X_test.reshape(-1, 28 * 28)

device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

in_feat = 28 * 28
out_feat = 10
net = LinearNet(in_feat, out_feat, hid_neurons=25)
net = net.to(device)

optimizer = torch.optim.Adam(net.parameters(), lr=3e-4)
criterion = torch.nn.CrossEntropyLoss()

num_epoch = 500
batch_size = 100

history_loss = {'train': [], 'test': []}
history_acc = {'train': [], 'test': []}

def generate_batch(X, y, batch_size):
    idx = torch.randperm(len(X))[:batch_size]
    return X[idx], y[idx]

iter_per_epoch = int(np.ceil(len(X_train) / batch_size))

X_test, y_test = X_test.to(device), y_test.to(device)

for epoch in range(num_epoch):

    total_epoch_loss = 0.0
    total_epoch_acc = 0.0
    for iter in range(iter_per_epoch):
        X_batch, y_batch = generate_batch(X_train, y_train, batch_size)
        X_batch = X_batch.to(device)
        y_batch = y_batch.to(device)

        output = net(X_batch)
        pred = torch.argmax(output, dim=1)
        batch_acc = (pred == y_batch).float().mean()
        total_epoch_acc += batch_acc

        loss = criterion(output, y_batch)
        total_epoch_loss += loss

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

    total_epoch_acc /= iter_per_epoch
    total_epoch_loss /= iter_per_epoch
    history_acc['train'].append(total_epoch_acc)
    history_loss['train'].append(total_epoch_loss)

    # validation
    with torch.no_grad():
        output = net(X_test)
        pred = torch.argmax(output, dim=1)
        acc = (pred == y_test).float().mean()
        loss = criterion(output, y_test)
        history_acc['test'].append(acc)
        history_loss['test'].append(loss)

        if epoch % 10 == 9: 
            print('Epoch {}, validation acc = {:.4f}'.format(epoch+1, acc))