Refer to tf.estimator.SessionRunHook when in tf 1.14 and beyond, otherwise use tf.train.SessionRunHook
+from copy import deepcopy
+
+from pkg_resources import parse_version
import tensorflow as tf
+
from wandb import util
+from wandb.apis.file_stream import Chunk
from wandb.data_types import history_dict_to_json
from wandb.tensorboard import *
import wandb
-from copy import deepcopy
-from wandb.apis.file_stream import Chunk
+
+tf_version = parse_version(tf.__version__)
+if tf_version >= parse_version('1.14'):
+ # In tf 1.14 and beyond, SessionRunHook is in the estimator package.
+ SessionRunHook = tf.estimator.SessionRunHook
+else:
+ SessionRunHook = tf.train.SessionRunHook
-class WandbHook(tf.train.SessionRunHook):
+class WandbHook(SessionRunHook):
def __init__(self, summary_op=None, steps_per_log=1000, history=None):
self._summary_op = summary_op
self._steps_per_log = steps_per_log
