experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Change the type of slice_size from int to auto to avoid overflow.

Browse Source 
PiperOrigin-RevId: 333791313
Change-Id: I4b395d87a50473bd2cac246d8144fe21ee3787ca
This commit is contained in:
A. Unique TensorFlower 2020-09-25 12:55:57 -07:00 committed by TensorFlower Gardener
parent 2ef29a9389
commit 94b6db8cc5
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/core/kernels/dynamic_stitch_op.cc
View File

@ -249,7 +249,8 @@ class DynamicStitchOpImplCPU : public DynamicStitchOpImplBase<T> {
// merged that aren't covered by an index in indices. What should we do?
if (first_dim_size > 0) {
auto merged_flat = merged->flat_outer_dims<T>();
const int slice_size = merged_flat.dimension(1);
// slice_size must not be stored as int for cases of tensors over 2GB.
const auto slice_size = merged_flat.dimension(1);
const size_t slice_bytes = slice_size * sizeof(T);
auto OnInputNumber = [&](int input_num) {
const Tensor& indices = indices_inputs[input_num];