Skip to content
Permalink
Browse files Browse the repository at this point in the history
Use PartialTensorShape instead of TensorShape.
`TensorShape` constructor throws a CHECK-fail if shape is partial/overflows which the other doesn't. We are only determining the number of elements in the shape and partial shape should be used as it returns negative number when needed.

PiperOrigin-RevId: 409205384
Change-Id: Ia56542ff9ec758f2c9ffc7e4dcc9fa7eecd86e7b
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Nov 11, 2021
1 parent 31d8f80 commit c2426bb
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion tensorflow/core/framework/attr_value_util.cc
Expand Up @@ -45,7 +45,7 @@ constexpr int kMaxTensorNestDepth = 100;
// not fully defined return -1.
int64_t TensorByteSize(const TensorProto& t) {
// num_elements returns -1 if shape is not fully defined.
int64_t num_elems = TensorShape(t.tensor_shape()).num_elements();
int64_t num_elems = PartialTensorShape(t.tensor_shape()).num_elements();
return num_elems < 0 ? -1 : num_elems * DataTypeSize(t.dtype());
}

Expand Down

0 comments on commit c2426bb

Please sign in to comment.