diff --git a/nn.py b/nn.py
index 7793afc..0e8be7a 100644
--- a/nn.py
+++ b/nn.py
@@ -16,14 +16,45 @@ tf.enable_eager_execution()
 df = pd.read_pickle('data_final.p')
 # print(df.head())
 
-# dataset = dataframe_to_dataset_biomes(df)
-dataset = dataframe_to_dataset_temp_precip(df)
+BATCH_SIZE = 15
+SHUFFLE_BUFFER_SIZE = 100
+LEARNING_RATE = 0.001
 
-i = 0
-for feature, target in dataset:
-    i += 1
-    if i > 10:
-        break
-    print('{} => {}'.format(feature, target))
+# dataset = dataframe_to_dataset_biomes(df)
+dataset_size, features, dataset = dataframe_to_dataset_temp_precip(df)
+print(dataset_size)
+dataset = dataset.shuffle(SHUFFLE_BUFFER_SIZE).batch(BATCH_SIZE).repeat()
+TRAIN_SIZE = dataset_size * 0.85
+TEST_SIZE = dataset_size - TRAIN_SIZE
+(training, test) = (dataset.take(TRAIN_SIZE), dataset.skip(TRAIN_SIZE))
+
+print(training.make_one_shot_iterator().get_next())
+
+model = keras.Sequential([
+    keras.layers.Dense(32, activation=tf.nn.relu, input_shape=[features]),
+    keras.layers.Dense(32, activation=tf.nn.relu),
+    keras.layers.Dense(2)
+])
+
+optimizer = tf.train.AdamOptimizer(LEARNING_RATE)
+
+model.compile(loss='mse',
+              optimizer=optimizer,
+              metrics=['mae'])
+
+model.summary()
+
+EPOCHS = 1000
+
+history = model.fit(
+    training,
+    epochs=EPOCHS,
+    verbose=1,
+    steps_per_epoch=int(dataset_size / BATCH_SIZE)
+)
+
+# i = 0
+# for feature, target in dataset:
+    # print('{} => {}'.format(feature, target))
 
 print(tf.__version__)
diff --git a/utils.py b/utils.py
index 000e571..99393ba 100644
--- a/utils.py
+++ b/utils.py
@@ -51,7 +51,7 @@ def dataframe_to_dataset_biomes(df):
     tf_inputs = tf.cast(normalize_ndarray(tf_inputs), tf.float32)
     tf_output = tf.cast(normalize_ndarray(tf_output), tf.int32)
 
-    return tf.data.Dataset.from_tensor_slices((tf_inputs, tf_output))
+    return int(tf_inputs.shape[0]), 5, tf.data.Dataset.from_tensor_slices((tf_inputs, tf_output))
 
 def dataframe_to_dataset_temp_precip(df):
     rows = df.shape[0]
@@ -81,5 +81,5 @@ def dataframe_to_dataset_temp_precip(df):
     tf_inputs = tf.cast(normalize_ndarray(tf_inputs), tf.float32)
     tf_output = tf.cast(normalize_ndarray(tf_output), tf.float32)
 
-    return tf.data.Dataset.from_tensor_slices((tf_inputs, tf_output))
+    return int(tf_inputs.shape[0]), 5, tf.data.Dataset.from_tensor_slices((tf_inputs, tf_output))