diff --git a/train_model_nn.py b/train_model_nn.py
index c3a3a52..89004d3 100644
--- a/train_model_nn.py
+++ b/train_model_nn.py
@@ -47,74 +47,50 @@ train_test_data = dict(zip(data_keys, data_vals))
 # %%
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
+TestNet2 = lambda: nn.Sequential(
+    nn.Linear(13, 32),
+    nn.LeakyReLU(),
+    nn.Linear(32, 16),
+    nn.Sigmoid(),
+    nn.Linear(16, 1),
+)
 
-class TestNet(nn.Module):
-    def __init__(self, n_input, n_hidden, n_output):
-        super().__init__()
-        self.hidden1 = nn.Linear(n_input, n_hidden)
-        self.hidden2 = nn.Linear(n_hidden, 16)
-        self.hidden3 = nn.Linear(16, 8)
-        self.hidden4 = nn.Linear(8, 4)
-        self.hidden5 = nn.Linear(4, 2)
-        self.predict = nn.Linear(2, n_output)
-
-    def forward(self, input):
-        out = self.hidden1(input)
-        out = F.relu(out)
-        out = F.normalize(out)
-        out = self.hidden2(out)
-        out = torch.sigmoid(out)
-        out = F.normalize(out)
-        out = self.hidden3(out)
-        out = F.relu(out)
-        out = F.normalize(out)
-        out = self.hidden4(out)
-        out = torch.sigmoid(out)
-        out = F.normalize(out)
-        out = self.hidden5(out)
-        out = F.relu(out)
-        out = F.normalize(out)
-        out = self.predict(out)
-        return out
-
-
-TestNet2 = nn.Sequential(nn.Linear(13, 32), nn.ReLU(), nn.Linear(32, 16),
-                         nn.Sigmoid(), nn.Linear(16, 1))
-
-# %%
 nets = {}
 
 X_train = torch.tensor(train_test_data["X_train"].values)
 y_train = train_test_data["y_train"]
 for target_col in y_train.columns:
-    y1 = torch.tensor(y_train[target_col].values)
+    y1 = torch.tensor(y_train[target_col].values).reshape(-1, 1)
     print(X.shape, y1.shape)
-    net = TestNet(13, 32, 1).double()
-    opti = torch.optim.SGD(net.parameters(), lr=0.1)
+    net = TestNet2().double()
+    opti = torch.optim.SGD(net.parameters(), lr=0.04)
     loss_func = nn.MSELoss()
-    for t in range(1000):
+    for t in range(10000):
         pred = net(X_train)
         loss = loss_func(pred, y1)
+        if t % 1000 == 0:
+            print(f'Epoch {t}, loss {loss}')
         opti.zero_grad()
         loss.backward()
         opti.step()
     nets[target_col] = net
     print(target_col)
+    break
 # %%
 X, ys = train_test_data['X_test'], train_test_data['y_test']
 evals = []
 from sklearn import metrics
-for target_col, net in nets.items():
-    y_hat = net(torch.tensor(X.values))  # fake
-    y_hat = y_hat.detach().numpy()
-    y = ys[target_col]  # real
-    print(y)
-    print(y_hat)
-    rmse = metrics.mean_squared_error(y, y_hat, squared=False)
-    r2 = metrics.r2_score(y, y_hat)
-    eval_dict = {'Error': target_col, 'RMSE': rmse, 'R^2': r2}
-    evals.append(eval_dict)
+with torch.no_grad():
+    for target_col, net in nets.items():
+        y_hat = net(torch.tensor(X.values))  # fake
+        y_hat = y_hat.detach().numpy()
+        y = ys[target_col]  # real
+        print(y)
+        print(y_hat)
+        rmse = metrics.mean_squared_error(y, y_hat, squared=False)
+        r2 = metrics.r2_score(y, y_hat)
+        eval_dict = {'Error': target_col, 'RMSE': rmse, 'R^2': r2}
+        evals.append(eval_dict)
 print(pd.DataFrame(evals))
 # %%