[keras] plot only on train end

keras.py

@@ -438,8 +438,8 @@ class PlotMulticlass(TFSummaryCallback):
     def on_train_begin(self, logs=None):
         self.make_input_plots()
 
-    def on_epoch_end(self, epoch, logs=None):
-        self.make_eval_plots(epoch)
+    def on_train_end(self, logs=None):
+        self.make_eval_plots()
 
     def make_input_plots(self):
         if self.plot_inputs: