diff --git a/assignment_1/src/t1.py b/assignment_1/src/t1.py
index f021bd7..23e14e7 100644
--- a/assignment_1/src/t1.py
+++ b/assignment_1/src/t1.py
@@ -60,24 +60,35 @@ X_val = X_val[:, 1:]
 
 mean = np.mean(X_train, axis=0)
 std = np.std(X_train, axis=0)
+y_mean = np.mean(y_train, axis=0)
+y_std = np.std(y_train, axis=0)
 
 X_train -= mean
 X_train /= std
+y_train -= y_mean
+y_train /= y_std
 
 X_val -= mean
 X_val /= std
+y_val -= y_mean
+y_val /= y_std
 
 network = Sequential()
-network.add(Dense(100, activation='relu'))
-network.add(Dense(100, activation='relu'))
+network.add(Dense(20, activation='relu'))
+network.add(Dense(20, activation='relu'))
 network.add(Dense(1, activation='sigmoid'))
 network.compile(optimizer='rmsprop', loss='mse', metrics=['mse'])
 
-network.fit(X_train, y_train, epochs=2000, verbose=1, batch_size=1000, validation_data=(X_val, y_val))
+callback = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=20)
+network.fit(X_train, y_train, epochs=2000, verbose=1, batch_size=1000, validation_data=(X_val, y_val), callbacks=[callback])
 
-X_val = A
-msq = np.mean((network.predict((X_val[:, 1:] - mean) / std) - y_val) ** 2)
+msq = np.mean(((network.predict(X_val) - y_val) * y_std + y_mean) ** 2)
 print(msq)
 
-msq = np.mean((network.predict((X_test[:, 1:] - mean) / std) - y_test) ** 2)
+X_test = X_test[:, 1:]
+X_test -= mean
+X_test /= std
+y_test -= y_mean
+y_test /= y_std
+msq = np.mean(((network.predict(X_test) - y_test) * y_std + y_mean) ** 2)
 print(msq)