smdebug_hook_pyt_loop.py

for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        
        # Use hook to record tensors
        hook.record_tensor_value(tensor_name="loss", tensor_value=loss)
        
        loss.backward()
        optimizer.step()