Move to batch end
@@ -244,14 +244,14 @@ class WandbCallback(keras.callbacks.Callback):
self._save_model(epoch)
def on_batch_begin(self, batch, logs=None):
+ pass
+
+ def on_batch_end(self, batch, logs=None):
if not self._graph_rendered:
# Couldn't do this in train_begin because keras may still not be compiled
wandb.run.summary['graph'] = wandb.Graph.from_keras(self.model)
self._graph_rendered = True
- def on_batch_end(self, batch, logs=None):
- pass
-
def on_train_begin(self, logs=None):
pass
