[Fix] Fix memory Leak in ZeroLike (#2664)

intel · Apr 9, 2024 · c97ab79 · c97ab79
1 parent a7c62a1
commit c97ab79
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/itex/core/kernels/gpu/constant_ops.cc b/itex/core/kernels/gpu/constant_ops.cc
@@ -123,11 +123,6 @@ class ZerosLikeOp<Device, Variant> : public OpKernel {
   explicit ZerosLikeOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
 
   void Compute(OpKernelContext* ctx) override {
-    const Tensor& input = ctx->input(0);
-
-    OP_REQUIRES(ctx, input.dims() == 0,
-                errors::InvalidArgument("ZerosLike non-scalar Tensor with "
-                                        "dtype=DT_VARIANT is not supported."));
     auto zeros_like_func = [](TF_OpKernelContext* tf_ctx, TF_Tensor* tf_input,
                               TF_Tensor* tf_out) {
       OpKernelContext ctx(tf_ctx);
@@ -147,6 +142,8 @@ class ZerosLikeOp<Device, Variant> : public OpKernel {
           break;
 #undef DTYPE_CASE
       }
+      TF_DeleteTensor(tf_input);
+      tf_input = nullptr;
     };
     TF_OpKernelContext* tf_ctx = ctx->Get();
     TF_Status* tf_status = TF_NewStatus();