Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix tf.raw_ops.GetSessionTensor and `tf.raw_ops.DeleteSessionTensor…
…` null pointer dereferences.

PiperOrigin-RevId: 368294154
Change-Id: Ie10f07a0a9a1c2b685e08153d48a0ca4b93f9fc9
  • Loading branch information
Amit Patankar authored and tensorflower-gardener committed Apr 13, 2021
1 parent eebb96c commit ff70c47
Showing 1 changed file with 10 additions and 4 deletions.
14 changes: 10 additions & 4 deletions tensorflow/core/kernels/session_ops.cc
Expand Up @@ -91,7 +91,6 @@ TF_CALL_NUMBER_TYPES(REGISTER_GPU_KERNEL);
REGISTER_GPU_KERNEL(bool);
#undef REGISTER_GPU_KERNEL


class GetSessionTensorOp : public OpKernel {
public:
explicit GetSessionTensorOp(OpKernelConstruction* context)
Expand All @@ -101,7 +100,11 @@ class GetSessionTensorOp : public OpKernel {
const Tensor& handle = ctx->input(0);
const string& name = handle.scalar<tstring>()();
Tensor val;
OP_REQUIRES_OK(ctx, ctx->session_state()->GetTensor(name, &val));
auto session_state = ctx->session_state();
OP_REQUIRES(ctx, session_state != nullptr,
errors::FailedPrecondition(
"GetSessionTensor called on null session state"));
OP_REQUIRES_OK(ctx, session_state->GetTensor(name, &val));
ctx->set_output(0, val);
}

Expand All @@ -122,7 +125,6 @@ TF_CALL_NUMBER_TYPES(REGISTER_GPU_KERNEL);
REGISTER_GPU_KERNEL(bool);
#undef REGISTER_GPU_KERNEL


class DeleteSessionTensorOp : public OpKernel {
public:
explicit DeleteSessionTensorOp(OpKernelConstruction* context)
Expand All @@ -131,7 +133,11 @@ class DeleteSessionTensorOp : public OpKernel {
void Compute(OpKernelContext* ctx) override {
const Tensor& handle = ctx->input(0);
const string& name = handle.scalar<tstring>()();
OP_REQUIRES_OK(ctx, ctx->session_state()->DeleteTensor(name));
auto session_state = ctx->session_state();
OP_REQUIRES(ctx, session_state != nullptr,
errors::FailedPrecondition(
"DeleteSessionTensor called on null session state"));
OP_REQUIRES_OK(ctx, session_state->DeleteTensor(name));
}

TF_DISALLOW_COPY_AND_ASSIGN(DeleteSessionTensorOp);
Expand Down

0 comments on commit ff70c47

Please sign in to comment.