docs: Update callback code snippets to include all imports needed for example#2283

ananthsub merged 1 commit intoNVIDIA-NeMo:mainNVIDIA-NeMo/Megatron-Bridge:mainfrom

ananthsub:callback-docs-updateananthsub/Megatron-Bridge:callback-docs-updateCopy head branch name to clipboard