diff options
Diffstat (limited to 'tensorflow/python/framework/tensor_util.py')
-rw-r--r-- | tensorflow/python/framework/tensor_util.py | 3 |
1 files changed, 3 insertions, 0 deletions
diff --git a/tensorflow/python/framework/tensor_util.py b/tensorflow/python/framework/tensor_util.py index d311cefcf4..ae4e73a363 100644 --- a/tensorflow/python/framework/tensor_util.py +++ b/tensorflow/python/framework/tensor_util.py @@ -361,6 +361,9 @@ def make_tensor_proto(values, dtype=None, shape=None): tensor_shape=tensor_shape.as_shape(shape).as_proto()) if is_same_size and numpy_dtype in _TENSOR_CONTENT_TYPES and shape_size > 1: + if nparray.size * nparray.itemsize >= (1 << 31): + raise ValueError( + "Cannot create a tensor proto whose content is larger than 2GB.") tensor_proto.tensor_content = nparray.tostring() return tensor_proto |