35 lines
1.3 KiB
Diff
35 lines
1.3 KiB
Diff
|
From 4f663d4b8f0bec1b48da6fa091a7d29609980fa4 Mon Sep 17 00:00:00 2001
|
||
|
|
From: Amit Patankar <amitpatankar@google.com>
|
||
|
|
Date: Mon, 8 Feb 2021 12:29:30 -0800
|
||
|
|
Subject: [PATCH] Allowlist certain data types to avoid a seg fault.
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 356326671
|
||
|
|
Change-Id: I23b65b52e93798cb5a6744632d31b0f88c6b6b31
|
||
|
|
---
|
||
|
|
tensorflow/core/kernels/immutable_constant_op.cc | 5 +++++
|
||
|
|
1 file changed, 5 insertions(+)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/kernels/immutable_constant_op.cc b/tensorflow/core/kernels/immutable_constant_op.cc
|
||
|
|
index 1cfbdb8277891..19aa865c1fbe4 100644
|
||
|
|
--- a/tensorflow/core/kernels/immutable_constant_op.cc
|
||
|
|
+++ b/tensorflow/core/kernels/immutable_constant_op.cc
|
||
|
|
@@ -17,6 +17,8 @@ limitations under the License.
|
||
|
|
|
||
|
|
#include <unordered_set>
|
||
|
|
|
||
|
|
+#include "tensorflow/core/framework/types.pb.h"
|
||
|
|
+
|
||
|
|
namespace tensorflow {
|
||
|
|
|
||
|
|
namespace {
|
||
|
|
@@ -86,6 +88,9 @@ ImmutableConstantOp::ImmutableConstantOp(OpKernelConstruction* context)
|
||
|
|
OP_REQUIRES_OK(context,
|
||
|
|
context->GetAttr(kMemoryRegionNameAttr, ®ion_name_));
|
||
|
|
OP_REQUIRES_OK(context, context->GetAttr(kDTypeAttr, &dtype_));
|
||
|
|
+ OP_REQUIRES(context, dtype_ != DT_RESOURCE && dtype_ != DT_VARIANT,
|
||
|
|
+ errors::InvalidArgument(
|
||
|
|
+ "Resource and variant dtypes are invalid for this op."));
|
||
|
|
OP_REQUIRES_OK(context, context->GetAttr(kShapeAttr, &shape_));
|
||
|
|
}
|
||
|
|
|