Ensure all-reduce statement for total loss after application of gradients during optimization.

video_prediction_tools/model_modules/video_prediction/models/vanilla_convLSTM_model.py

@@ -153,11 +153,12 @@ class VanillaConvLstmVideoPredictionModel(object):
         # If monitoredTrainingSession is used, one must make use of hvd.BroadcastGlobalVariableHook(0),
         # If not, hvd.broadcast_global_variables after initialization of gloabl variables is suifficient
         # self.hooks = [hvd.BroadcastGlobalVariablesHook(0)]
-        with tf.control_dependencies([grads]):
-            self.total_loss = hvd.allreduce(self.total_loss) 
         self.bcast_op = hvd.broadcast_global_variables(0)
         # Apply gradient during training iteration
         self.train_op = opt.apply_gradients(grads_and_vars=grads, global_step=self.global_step)
+        # average loss
+        with tf.control_dependencies([self.train_op]):
+            self.total_loss = hvd.allreduce(self.total_loss) 
 
         self.outputs = {}
         self.outputs["gen_images"] = self.x_hat