Allowlist certain data types to avoid a seg fault.
PiperOrigin-RevId: 356326671 Change-Id: I23b65b52e93798cb5a6744632d31b0f88c6b6b31
This commit is contained in:
parent
f0e867da82
commit
4f663d4b8f
@ -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_));
|
||||
}
|
||||
|
||||
|
Loading…
x
Reference in New Issue
Block a user