From 78582dbecd0c9b9ac30ee7cdf5ffbfdb8dd2863c Mon Sep 17 00:00:00 2001 From: chenfeiyu Date: Mon, 24 Feb 2020 06:54:57 +0000 Subject: [PATCH] make DataIterator compatible for python 2 --- parakeet/data/datacargo.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/parakeet/data/datacargo.py b/parakeet/data/datacargo.py index 40bcb02..8c9a3b2 100644 --- a/parakeet/data/datacargo.py +++ b/parakeet/data/datacargo.py @@ -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)