Merge branch 'michael_issue#132_bugfix_loss_tracking' into bing_issue#131_extract_g1_loss

video_prediction_tools/main_scripts/main_train_models.py

@@ -310,12 +310,14 @@ class TrainModel(object):
                 # Run training data
                 self.create_fetches_for_train()             # In addition to the loss, we fetch the optimizer
                 self.results = sess.run(self.fetches)       # ...and run it here!
-                train_losses.append(self.results[self.saver_loss])
+                # Note: For SAVP, the obtained loss is a list where the first element is of interest, for convLSTM,
+                # it's just a number. Thus, with list(<losses>)[0], we can handle both
+                train_losses.append(list(self.results[self.saver_loss])[0])
                 # run and fetch losses for validation data
                 val_handle_eval = sess.run(self.val_handle)
                 self.create_fetches_for_val()
                 self.val_results = sess.run(self.val_fetches, feed_dict={self.train_handle: val_handle_eval})
-                val_losses.append(self.val_results[self.saver_loss])
+                val_losses.append(list(self.val_results[self.saver_loss])[0])
                 self.write_to_summary()
                 self.print_results(step, self.results)
                 # track iteration time