Override the function for returning next batch.

def next(self):
        """Override the function for returning next batch."""
        batch_size = self.batch_size
        c, h, w = self.data_shape
        # if last batch data is rolled over
        if self._cache_data is not None:
            # check both the data and label have values
            assert self._cache_label is not None, "_cache_label didn't have values"
            assert self._cache_idx is not None, "_cache_idx didn't have values"
            batch_data = self._cache_data
            batch_label = self._cache_label
            i = self._cache_idx
        else:
            batch_data = nd.zeros((batch_size, c, h, w))
            batch_label = nd.empty(self.provide_label[0][1])
            batch_label[:] = -1
            i = self._batchify(batch_data, batch_label)
        # calculate the padding
        pad = batch_size - i
        # handle padding for the last batch
        if pad != 0:
            if self.last_batch_handle == 'discard':
                raise StopIteration
            # if the option is 'roll_over', throw StopIteration and cache the data
            elif self.last_batch_handle == 'roll_over' and \
                self._cache_data is None:
                self._cache_data = batch_data
                self._cache_label = batch_label
                self._cache_idx = i
                raise StopIteration
            else:
                _ = self._batchify(batch_data, batch_label, i)
                if self.last_batch_handle == 'pad':
                    self._allow_read = False
                else:
                    self._cache_data = None
                    self._cache_label = None
                    self._cache_idx = None

        return io.DataBatch([batch_data], [batch_label], pad=pad)