smdebug_hook_pyt.py

import smdebug.pytorch as smd

net = get_network()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

job_name = 'pytorch-debug-job'
hook = smd.Hook(out_dir=f'./smd_outputs/{job_name}',
               save_config=smd.SaveConfig(save_interval=10),
               include_collections=['gradients', 'biases'])

hook.register_module(net)
hook.register_loss(criterion)