def METHOD_NAME(output_s3_uri):
    # With the following SaveConfig, we will save tensors for steps 1, 2 and 3
    # (indexing starts with 0).
    save_config = SaveConfig(save_steps=[1, 2, 3])
    # Create a hook that logs weights, biases and gradients while training the model.
    hook = Hook(
        out_dir=output_s3_uri,
        save_config=save_config,
        include_collections=["weights", "gradients"],
    )
    return hook