criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
train_loss_mlp = []
test_loss_mlp = []

epochs = 250
for epoch in range(epochs):
    optimizer.zero_grad()
    output = model(dataset['train_input']).squeeze()
    loss = criterion(output, dataset['train_label'])
    loss.backward()
    optimizer.step()
    train_loss_mlp.append(loss.item()**0.5)

    # Test the model
    model.eval()
    with torch.no_grad():
        output = model(dataset['test_input']).squeeze()
        loss = criterion(output, dataset['test_label'])
        test_loss_mlp.append(loss.item()**0.5)

    print(f'Epoch {epoch+1}/{epochs}, Train Loss: {train_loss_mlp[-1]:.2f}, Test Loss: {test_loss_mlp[-1]:.2f}', end='r')