diff --git a/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java b/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java
index 3c4e7762e3..53aa0fa267 100644
--- a/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java
+++ b/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java
@@ -92,6 +92,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLTask
             ActionListener<MLModel> modelActionListener = new ActionListener<>() {
                 @Override
                 public void onResponse(MLModel mlModel) {
+                    context.restore();
                     modelCacheHelper.setModelInfo(modelId, mlModel);
                     FunctionName functionName = mlModel.getAlgorithm();
                     mlPredictionTaskRequest.getMlInput().setAlgorithm(functionName);