diff --git a/.gitignore b/.gitignore
index 82ed5e9..aa22a7f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -114,3 +114,6 @@ tf_pywatts_model/
 
 # Tensorboard
 pywatts/tensorboard
+
+# Figures
+figures/
diff --git a/pywatts/eval_training.py b/pywatts/eval_training.py
index 22ddd42..ab95090 100644
--- a/pywatts/eval_training.py
+++ b/pywatts/eval_training.py
@@ -54,15 +54,21 @@ if TRAIN:
                 pp.plot(pred_result, 'black')
 
             pp.plot(prediction, color=color_gradient_base)
+            pp.savefig(FIGURE_OUTPUT_DIR+'{}.pdf'.format(q), orientation='landscape')
 
         color_gradient_base = tuple([sum(x) for x in zip(color_gradient_base, color_step_width)])
 
     for i in range(NUM_QUERIES):
-        pp.figure(i)
-        pp.savefig(FIGURE_OUTPUT_DIR+'{}.pdf'.format(i), orientation='landscape')
+        pp.close(i)
 
     if PLOT:
         # Plot training success rate (with 'average loss')
-        pywatts.routines.plot_training(train_eval)
+        loss = []
+        for e in train_eval:
+            loss.append(e['average_loss'])
+
+        pp.plot(loss)
+        # Needed for execution in PyCharm
+        pp.show()
 
 exit()
diff --git a/pywatts/test_kcross_train.py b/pywatts/test_kcross_train.py
index 6201fa2..4e29d0a 100644
--- a/pywatts/test_kcross_train.py
+++ b/pywatts/test_kcross_train.py
@@ -35,7 +35,7 @@ if TRAIN:
 
 if PLOT:
     # Plot training success rate (with 'average loss')
-    pywatts.main.plot_training(train_eval)
+    pywatts.routines.plot_training(train_eval)
 
 
 exit()