Skip to content

Commit 273f489

Browse files
committed
update recordio unittest
1 parent 75a2e50 commit 273f489

File tree

2 files changed

+11
-3
lines changed

2 files changed

+11
-3
lines changed

paddle/fluid/framework/reader.cc

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@ FileReader::FileReader(const std::vector<DDim> &dims) : dims_(dims) {}
2222

2323
void FileReader::ReadNext(std::vector<LoDTensor> *out) {
2424
ReadNextImpl(out);
25-
PADDLE_ENFORCE_EQ(out->size(), dims_.size());
25+
if (out->empty()) {
26+
return;
27+
}
2628
for (size_t i = 0; i < dims_.size(); ++i) {
2729
auto &actual = out->at(i).dims();
2830
auto &expect = dims_[i];

python/paddle/fluid/tests/unittests/test_recordio_reader.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,8 +65,14 @@ def test_main(self, decorator_callback=None):
6565

6666
# train a pass
6767
batch_id = 0
68-
while not data_file.eof():
69-
tmp, = exe.run(fetch_list=[avg_loss])
68+
while True:
69+
ex = None
70+
try:
71+
tmp, = exe.run(fetch_list=[avg_loss])
72+
except fluid.core.EnforceNotMet as ex:
73+
self.assertIn("There is no next data.", ex.message)
74+
break
75+
7076
avg_loss_np.append(tmp)
7177
batch_id += 1
7278
data_file.reset()

0 commit comments

Comments
 (0)