check status of allocate_temp (deepmodeling#2797)

It's the same as deepmodeling#2782, but one line was missing in that PR. Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Sep 8, 2023 · c6829bc · c6829bc
1 parent c6dea0c
commit c6829bc
Showing 1 changed file with 5 additions and 2 deletions.
diff --git 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);