make DataIterator compatible for python 2

This commit is contained in:
chenfeiyu 2020-02-24 06:54:57 +00:00
parent 25883dcd3e
commit 78582dbecd
1 changed files with 6 additions and 1 deletions

View File

@ -1,3 +1,4 @@
import six
from .sampler import SequentialSampler, RandomSampler, BatchSampler
@ -84,7 +85,11 @@ class DataIterator(object):
return minibatch
def _next_index(self):
return next(self._sampler_iter)
if six.PY3:
return next(self._sampler_iter)
else:
# six.PY2
return self._sampler_iter.next()
def __len__(self):
return len(self._index_sampler)