cp: docs: Update callback code snippets to include all imports needed for example (2283) into r0.3.0#2291

ko3n1g merged 1 commit into

r0.3.0NVIDIA-NeMo/Megatron-Bridge:r0.3.0from

cherry-pick-2283-r0.3.0NVIDIA-NeMo/Megatron-Bridge:cherry-pick-2283-r0.3.0Copy head branch name to clipboard