diff --git a/video_prediction_savp/scripts/train_dummy.py b/video_prediction_savp/scripts/train_dummy.py
index 1fb401955c39be4807cf7747e43ed660941cb925..805e81d79707338665f0fe2b1fad3b212359c233 100644
--- a/video_prediction_savp/scripts/train_dummy.py
+++ b/video_prediction_savp/scripts/train_dummy.py
@@ -160,11 +160,11 @@ def make_dataset_iterator(train_dataset, val_dataset, batch_size ):
 
 
 def plot_train(train_losses,val_losses,output_dir):
-    epochs = list(range(len(train_losses))) 
-    plt.plot(epochs, train_losses, 'g', label='Training loss')
-    plt.plot(epochs, val_losses, 'b', label='validation loss')
+    iterations = list(range(len(train_losses))) 
+    plt.plot(iterations, train_losses, 'g', label='Training loss')
+    plt.plot(iterations, val_losses, 'b', label='validation loss')
     plt.title('Training and Validation loss')
-    plt.xlabel('Epochs')
+    plt.xlabel('Iterations')
     plt.ylabel('Loss')
     plt.legend()
     plt.savefig(os.path.join(output_dir,'plot_train.png'))
diff --git a/video_prediction_savp/video_prediction/models/vanilla_convLSTM_model.py b/video_prediction_savp/video_prediction/models/vanilla_convLSTM_model.py
index 7e3fec28dc28c78b8203e1924f17489af8f5075e..557280d5d7169b9212844d4739ce3e0c2df7190b 100644
--- a/video_prediction_savp/video_prediction/models/vanilla_convLSTM_model.py
+++ b/video_prediction_savp/video_prediction/models/vanilla_convLSTM_model.py
@@ -69,6 +69,7 @@ class VanillaConvLstmVideoPredictionModel(BaseVideoPredictionModel):
         #print("self.x_hat_context_frames,",self.x_hat_context_frames)
         #self.context_frames_loss = tf.reduce_mean(
         #    tf.square(self.x[:, :self.context_frames, :, :, 0] - self.x_hat_context_frames[:, :, :, :, 0]))
+        # This is the loss function (RMSE):
         self.total_loss = tf.reduce_mean(
             tf.square(self.x[:, self.context_frames:, :, :, 0] - self.x_hat_context_frames[:, (self.context_frames-1):-1, :, :, 0]))