Set mode on each batch (aws#65)

Author: rahul003

smdebug/tensorflow/keras.py

@@ -479,6 +479,11 @@ def _on_any_batch_begin(self, batch, mode, logs=None):
         if self._is_not_supported():
             return
 
+        # set mode for each batch as when users run model.fit() and pass validation data
+        # through the optional argument, then mode_begin is not called for the training steps
+        # after first evaluation during training
+        self.set_mode(mode)
+
         self._close_writers()
         self._increment_step()