Mercurial > ift6266
diff datasets/dsetiter.py @ 189:0d0677773533
Fix bug where there would be a bunch of 0-length batches at the end under certain circumstances.
author | Arnaud Bergeron <abergeron@gmail.com> |
---|---|
date | Mon, 01 Mar 2010 17:06:49 -0500 |
parents | 76bc047df5ee |
children | 1adfafdc3d57 |
line wrap: on
line diff
--- a/datasets/dsetiter.py Mon Mar 01 11:46:39 2010 -0500 +++ b/datasets/dsetiter.py Mon Mar 01 17:06:49 2010 -0500 @@ -123,6 +123,7 @@ tmpbuf = self.curfile.read(self.bufsize - len(buf)) buf = numpy.row_stack((buf, tmpbuf)) + self.cursize = len(buf) self.buffer = buf self.curpos = 0 @@ -144,9 +145,17 @@ Traceback (most recent call last): ... StopIteration - + >>> d = DataIterator([DummyFile(13)], 10, 50) + >>> len(d.next()) + 10 + >>> len(d.next()) + 3 + >>> d.next() + Traceback (most recent call last): + ... + StopIteration """ - if self.curpos >= self.bufsize: + if self.curpos >= self.cursize: self._fill_buf() res = self.buffer[self.curpos:self.curpos+self.batchsize] self.curpos += self.batchsize