smdebug_hook.py

import smdebug.tensorflow as smd
job_name = 'tf-debug-job'

hook = smd.KerasHook(out_dir=f'./smd_outputs/{job_name}',
                    tensorboard_dir=f'./tb_logs/{job_name}',
                    save_config=smd.SaveConfig(save_interval=1),
                    include_collections=['gradients', 'biases'])
                    
opt = tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.9, name='SGD')
opt = hook.wrap_optimizer(opt)