diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py
index 295fb633ff..e9bae28732 100644
--- a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py
+++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py
@@ -202,7 +202,7 @@ def _find_dependencies_in_step_arguments(
                 else:
                     dependencies.add(self._get_step_name_from_str(referenced_step, step_map))
 
-            from sagemaker.core.workflow.function_step import DelayedReturn
+            from sagemaker.mlops.workflow.function_step import DelayedReturn
 
             # TODO: we can remove the if-elif once move the validators to JsonGet constructor
             if isinstance(pipeline_variable, JsonGet):