diff --git a/pywatts/main.py b/pywatts/main.py
index 04815ed..b509c22 100644
--- a/pywatts/main.py
+++ b/pywatts/main.py
@@ -4,15 +4,16 @@ import matplotlib.pyplot as pp
 import pywatts.neural
 from sklearn.metrics import explained_variance_score, mean_absolute_error, median_absolute_error
 import pandas
+from random import randint
 
 from sklearn.model_selection import train_test_split
 
 
 df = pywatts.db.rows_to_df(list(range(1, 50)))
-X = df[[col for col in df.columns if col != 'dc']]
+X = df
 y = df['dc']
 
-X_train, X_tmp, y_train, y_tmp = train_test_split(X, y, test_size=0.2, random_state=23)
+X_train, X_tmp, y_train, y_tmp = train_test_split(X, y, test_size=0.2, random_state=34)
 
 X_test, X_val, y_test, y_val = train_test_split(X_tmp, y_tmp, test_size=0.5, random_state=23)
 
@@ -20,20 +21,34 @@ feature_cols = [tf.feature_column.numeric_column(col) for col in X.columns]
 n = pywatts.neural.Net(feature_cols=feature_cols)
 
 
-def input_data(json_str):
+def train_split(data, size):
+    X_values = {'dc': [], 'temp': [], 'wind': []}
+    y_values = []
+    for i in range(size):
+        rnd_idx = randint(0, data.size / data.shape[1] - 337)
+
+        X_values['dc'].extend(data['dc'][rnd_idx:rnd_idx + 336])
+        X_values['temp'].extend(data['temp'][rnd_idx:rnd_idx + 336])
+        X_values['wind'].extend(data['wind'][rnd_idx:rnd_idx + 336])
+        y_values.append(data['dc'][rnd_idx + 337])
+
+    return pandas.DataFrame.from_dict(X_values), pandas.DataFrame.from_dict({'dc': y_values})
+
+
+def input_data(json_str, idx=0):
     tmp_df = pandas.read_json(json_str)
 
     return pandas.DataFrame.from_dict(
-        {'dc': [x for l in tmp_df['dc'] for x in l],
-         'temp': [x for l in tmp_df['temp'] for x in l],
-         'wind': [x for l in tmp_df['wind'] for x in l]}
+        {'dc': tmp_df['dc'][idx],
+         'temp': tmp_df['temp'][idx],
+         'wind': tmp_df['wind'][idx]}
     )
 
 
 def train(steps=100):
     evaluation = []
     for i in range(steps):
-        n.train(X_train, y_train, steps=400)
+        n.train(X_train, y_train, steps=100)
         evaluation.append(n.evaluate(X_val, y_val))
         print("Training %s of %s" % ((i+1), steps))
     return evaluation
@@ -42,7 +57,7 @@ def train(steps=100):
 def plot_training(evaluation):
     loss = []
     for e in evaluation:
-        loss.append(e['loss'])
+        loss.append(e['average_loss'])
     pp.plot(loss)
 
 
@@ -52,11 +67,11 @@ def predict(X_pred):
     return predictions
 
 
-def eval_prediction(prediction):
+def eval_prediction(prediction, result):
     print("The Explained Variance: %.2f" % explained_variance_score(
-        y_test, prediction))
+        result, prediction))
     print("The Mean Absolute Error: %.2f volt dc" % mean_absolute_error(
-        y_test, prediction))
+        result, prediction))
     print("The Median Absolute Error: %.2f volt dc" % median_absolute_error(
-        y_test, prediction))
+        result, prediction))
 
diff --git a/pywatts/neural.py b/pywatts/neural.py
index 673adc9..c1fbfc5 100644
--- a/pywatts/neural.py
+++ b/pywatts/neural.py
@@ -1,7 +1,7 @@
 import tensorflow as tf
 
 
-def pywatts_input_fn(X, y=None, num_epochs=None, shuffle=True, batch_size=400):
+def pywatts_input_fn(X, y=None, num_epochs=None, shuffle=True, batch_size=366):
     return tf.estimator.inputs.pandas_input_fn(x=X,
                                                y=y,
                                                num_epochs=num_epochs,
@@ -15,11 +15,11 @@ class Net:
 
     def __init__(self, feature_cols=__feature_cols):
         self.__regressor = tf.estimator.DNNRegressor(feature_columns=feature_cols,
-                                                     hidden_units=[50, 50],
+                                                     hidden_units=[2],
                                                      model_dir='tf_pywatts_model')
 
     def train(self, training_data, training_results, steps):
-        self.__regressor.train(input_fn=pywatts_input_fn(training_data, y=training_results, num_epochs=None, shuffle=True), steps=steps)
+        self.__regressor.train(input_fn=pywatts_input_fn(training_data, y=training_results, num_epochs=None, shuffle=True, batch_size=336), steps=steps)
 
     def evaluate(self, eval_data, eval_results):
         return self.__regressor.evaluate(input_fn=pywatts_input_fn(eval_data, y=eval_results, num_epochs=1, shuffle=False), steps=1)