Commit 766a1e00 authored by Ruoxin Sang's avatar Ruoxin Sang Committed by TensorFlower Gardener
Browse files

Make DS InputIterator get_next() API behave correctly in the last partial... 

Make DS InputIterator get_next() API behave correctly in the last partial batch case with multiple devices/workers.

Now there are 3 cases for get_next():
1. Each replica gets a full batch, the behavior is the same as before.
2. Some replicas get full batches, some get partial batches, and some get no data. get_next() will return a list with tensors from all replicas which include partial batch data and tensors with batch dimension 0 representing no data.
3. If there is no data in any replicas, an OutOfRange error will be triggered.

PiperOrigin-RevId: 238491718
parent 0094a3fd
Show whitespace changes
Inline Side-by-side
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment