Fix corner case that batch size is larger than dataset length in coco record dataset.
@@ -98,16 +98,23 @@ class TFDataDataLoader(BaseDataLoader):
# pylint: disable=no-name-in-module
from tensorflow.python.framework.errors_impl import OutOfRangeError
while True:
- try:
if not try_single_batch:
+ try:
outputs = data_sess.run(iter_tensors)
yield outputs
+ except OutOfRangeError:
+ data_sess.close()
+ return
else:
- outputs = [squeeze_output(data_sess.run(iter_tensors)) \
- for i in range(0, batch_size) ]
+ try:
+ outputs = []
+ for i in range(0, batch_size):
+ outputs.append(squeeze_output(data_sess.run(iter_tensors)))
outputs = default_collate(outputs)
yield outputs
except OutOfRangeError:
+ outputs = default_collate(outputs)
+ yield outputs
data_sess.close()
return
