From e852203f88373c393478260e639f37fc27c7576d Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Thu, 7 Sep 2023 16:55:21 -0400
Subject: [PATCH] check status of allocate_temp

Same as #2782, but one line was missing is that PR.

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 source/op/prod_env_mat_multi_device.cc | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/source/op/prod_env_mat_multi_device.cc b/source/op/prod_env_mat_multi_device.cc
index d423e8a108..a8882fb5f4 100644
--- a/source/op/prod_env_mat_multi_device.cc
+++ b/source/op/prod_env_mat_multi_device.cc
@@ -1549,8 +1549,11 @@ static int _norm_copy_coord_gpu(OpKernelContext* context,
     // Tensor cpy_temp;
     TensorShape cpy_shape;
     cpy_shape.AddDim(mem_cpy * 3);
-    context->allocate_temp(DataTypeToEnum<FPTYPE>::value, cpy_shape,
-                           tensor_list + 3);
+    status = context->allocate_temp(DataTypeToEnum<FPTYPE>::value, cpy_shape,
+                                    tensor_list + 3);
+    if (!status.ok()) {
+      return false;
+    }
     // Tensor t_temp;
     TensorShape t_shape;
     t_shape.AddDim(mem_cpy * 2);