Fix callback for tensorflow.keras
@@ -254,7 +254,10 @@ class WandbCallback(keras.callbacks.Callback):
pass
def on_train_batch_end(self, batch, logs=None):
- pass
+ if not self._graph_rendered:
+ # Couldn't do this in train_begin because keras may still not be built
+ wandb.run.summary['graph'] = wandb.Graph.from_keras(self.model)
+ self._graph_rendered = True
def on_test_begin(self, logs=None):
pass
