Fix broken keras image logging with tf1
@@ -56,7 +56,7 @@ def patch_tf_keras():
from tensorflow.python.keras.engine import training_arrays
from tensorflow.python.keras.engine import training_generator
- training_v2 = wandb.util.import_module('tensorflow.python.keras.engine.training_v2')
+ training_v2 = wandb.util.get_module('tensorflow.python.keras.engine.training_v2')
old_arrays = training_arrays.fit_loop
old_generator = training_generator.fit_generator
if training_v2:
