diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py
index b948976efd..1a53787002 100644
--- a/examples/inference/gpt/gpt_dynamic_inference.py
+++ b/examples/inference/gpt/gpt_dynamic_inference.py
@@ -30,7 +30,7 @@
     ContextOverflowError,
     DynamicInferenceContext,
 )
-from megatron.core.inference.context.attention_context.mamba_metadata import (
+from megatron.core.inference.contexts.attention_context.mamba_metadata import (
     MambaInferenceStateConfig,
 )
 from megatron.core.inference.engines import DynamicInferenceEngine, EngineSuspendedError