diff --git a/create-dataset.py b/create-dataset.py
index 049ac50..64744c0 100644
--- a/create-dataset.py
+++ b/create-dataset.py
@@ -1,6 +1,7 @@
 import pandas as pd
 import os
 from sklearn.model_selection import train_test_split
+import train
 
 CUTOFF = int(os.environ['CUTOFF'])
 adults = pd.read_csv('adult.csv')
@@ -20,3 +21,5 @@ X_test.to_csv('X_test.csv', index=False)
 Y_test.to_csv('Y_test.csv', index=False)
 Y_train.to_csv('Y_train.csv', index=False)
 Y_dev.to_csv('Y_dev.csv', index=False)
+
+train.main()
\ No newline at end of file
diff --git a/train.py b/train.py
index 79987eb..6594e7f 100644
--- a/train.py
+++ b/train.py
@@ -2,23 +2,25 @@ import pandas as pd
 import tensorflow
 from keras.applications.densenet import layers
 
-train_data_x = pd.read_csv('./X_train.csv')
 
-adults_train = train_data_x.copy()
-adults_predict = train_data_x.pop('age')
-normalize = layers.Normalization()
-normalize.adapt(adults_train)
+def main():
+    train_data_x = pd.read_csv('./X_train.csv')
 
-adult_model = tensorflow.keras.Sequential([
-  normalize,
-  layers.Dense(64),
-  layers.Dense(1)
-])
+    adults_train = train_data_x.copy()
+    adults_predict = train_data_x.pop('age')
+    normalize = layers.Normalization()
+    normalize.adapt(adults_train)
 
-adult_model.compile(
-    loss=tensorflow.keras.losses.MeanSquaredError(),
-    optimizer=tensorflow.keras.optimizers.Adam())
+    adult_model = tensorflow.keras.Sequential([
+        normalize,
+        layers.Dense(64),
+        layers.Dense(1)
+    ])
 
-adult_model.fit(adults_train, adults_predict, epochs=500)
+    adult_model.compile(
+        loss=tensorflow.keras.losses.MeanSquaredError(),
+        optimizer=tensorflow.keras.optimizers.Adam())
 
-adult_model.save('model')
\ No newline at end of file
+    adult_model.fit(adults_train, adults_predict, epochs=500)
+
+    adult_model.save('model')