diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/InferTrainedModelDeploymentAction.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/InferTrainedModelDeploymentAction.java
index f18d10ef3cacf..f384c5c708180 100644
--- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/InferTrainedModelDeploymentAction.java
+++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/InferTrainedModelDeploymentAction.java
@@ -129,5 +129,11 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             result.toXContent(builder, params);
             return builder;
         }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            super.writeTo(out);
+            result.writeTo(out);
+        }
     }
 }
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/deployment/PyTorchResult.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/deployment/PyTorchResult.java
index 3d776ee7a41c2..e5115357fbde6 100644
--- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/deployment/PyTorchResult.java
+++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/deployment/PyTorchResult.java
@@ -46,11 +46,7 @@ public class PyTorchResult implements ToXContentObject, Writeable {
                 double[][] primitiveDoubles = new double[listOfListOfDoubles.size()][];
                 for (int i = 0; i < listOfListOfDoubles.size(); i++) {
                     List<Double> row = listOfListOfDoubles.get(i);
-                    double[] primitiveRow = new double[row.size()];
-                    for (int j = 0; j < row.size(); j++) {
-                        primitiveRow[j] = row.get(j);
-                    }
-                    primitiveDoubles[i] = primitiveRow;
+                    primitiveDoubles[i] = row.stream().mapToDouble(d -> d).toArray();
                 }
                 return primitiveDoubles;
             },
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/deployment/DeploymentManager.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/deployment/DeploymentManager.java
index 35a90bb092006..176a463e50c07 100644
--- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/deployment/DeploymentManager.java
+++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/deployment/DeploymentManager.java
@@ -11,11 +11,13 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.apache.lucene.util.SetOnce;
+import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.get.GetRequest;
 import org.elasticsearch.action.get.GetResponse;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.common.unit.TimeValue;
+import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.core.ml.inference.deployment.PyTorchResult;
 import org.elasticsearch.xpack.core.ml.inference.deployment.TrainedModelDeploymentState;
@@ -113,7 +115,8 @@ private void waitForResult(ProcessContext processContext, String requestId, Acti
             TimeValue timeout = TimeValue.timeValueSeconds(5);
             PyTorchResult pyTorchResult = processContext.resultProcessor.waitForResult(requestId, timeout);
             if (pyTorchResult == null) {
-                listener.onFailure(ExceptionsHelper.serverError("no result was produced within timeout value [{}]", timeout));
+                listener.onFailure(new ElasticsearchStatusException("timeout [{}] waiting for inference result",
+                    RestStatus.TOO_MANY_REQUESTS, timeout));
             } else {
                 listener.onResponse(pyTorchResult);
             }