StopIteration is not working in Custom Iterable Dataloader

I have written a custom Iterable Dataloader for my project. While testing it for a single epoch, it goes through the whole data, then resets the counter to zero and starts iterating again bypassing the “raise StopIteration” command. I am a bit confused why it is working this way. Any help would be highly appreciated.

The next() function is given below

def __next__(self):
        if(self._counter < len(self.index)):
            self.sample = self._data[self.index[self._counter] : self.index[self._counter] + self._chunk_size, :]
            self.label = self._labels[self.index[self._counter]]
            
            if self._mean_std_normalize:
                self.sample = (self.sample - 17041.3242) / (131.5367)
            
            print(self._counter)
            self._counter += 1
            return self.sample, self.label
        else:
             self._counter = 0
             raise StopIteration

While testing it using a DataLoader, it works completely fine with a batch_size = 5, and runs in infinite loop for other batch_sizes. The testing script is

train_data = IterDataset(root_path = 'D:/idrad', folder = 'train',
                          chunk_size = 30, mean_std_normalize = True, shuffle = True)
train_loader = DataLoader(dataset = train_data, batch_size = 10)

epochs = 1
for i in range(epochs):
    for idx, (md, target) in enumerate(train_loader):
        pass

@ptrblck kindly help me to sort out this issue.

Other members are also encouraged to help me.

Thanking in advance

This solved my problem.

def __next__(self):
        if(self._counter < len(self.index)):
            self.sample = self._data[self.index[self._counter] : self.index[self._counter] + self._chunk_size, :]
            self.label = self._labels[self.index[self._counter]]
            
            if self._mean_std_normalize:
                self.sample = (self.sample - 17041.3242) / (131.5367)
            
            print(self._counter)
            self._counter += 1
            return self.sample, self.label
        else:
             raise StopIteration

and moving self._counter = 0 to __iter__ method. Now, my __iter__ method is

def __iter__(self):
    self._counter = 0
    return self