Doc loading checkpoint

[robmarkcole]

Added this cell, might be useful to include:

# load previously trained model
checkpoint_path = "../output6/Unet_ep100/chkpt/Unet_final_state.pth"
checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu')) # or use GPU

# Remove 'module.' prefix if present (for nn.DataParallel compatibility)
new_state_dict = {k.replace('module.', ''): v for k, v in checkpoint.items()}
model.load_state_dict(new_state_dict)