A real training set rarely fits neatly into the shape a model's forward pass expects, delivered in exactly the right batch size, in exactly the right order, exactly once per epoch. It usually lives as a big pile of samples (rows in a CSV, image files in a folder, entries in a database), and turning that pile into a stream of correctly-shaped, correctly-shuffled batches, every single epoch, is genuinely fiddly bookkeeping: how many full batches fit, what to do with the leftover samples that don't fill a full batch, how to shuffle WITHOUT accidentally mismatching a feature with the wrong target.
Splitting this into two separate pieces, a Dataset (which only knows how to answer "how many samples do you have" and "give me sample number i") and a DataLoader (which only knows how to turn a Dataset into a stream of shuffled batches), is what lets the SAME DataLoader logic work for a tiny in-memory array of 100 samples and a dataset that streams gigabytes of images off disk: DataLoader never needs to know or care how a Dataset actually stores its data, only that it can answer len(dataset) and dataset[i].
Implement ArrayDataset.__len__ and __getitem__ (the simplest possible Dataset, wrapping two parallel arrays), and DataLoader.__iter__ and __len__. DataLoader.__iter__ should build an index array [0, 1, ..., len(dataset)-1], shuffle it in place (using np.random.RandomState(self.seed) for reproducibility) if self.shuffle is True, then walk through it in chunks of self.batch_size, using each chunk of indices to gather and stack the corresponding samples from self.dataset, yielding one (batch_features, batch_targets) pair per chunk.
ArrayDataset.__getitem__(idx) returns a (feature, target) pair for that single index, not a batch.DataLoader.__iter__ must yield every sample in the dataset EXACTLY once per full iteration, the last batch may be smaller than batch_size if the dataset size doesn't divide evenly, but no sample is skipped or duplicated.shuffle=True, use np.random.RandomState(self.seed) specifically (not the global np.random state), so results are reproducible given the same seed.DataLoader.__len__ returns the number of batches a full pass produces (ceil(len(dataset) / batch_size)), not the number of samples.__len__ is just len(self.features); __getitem__(idx) is return self.features[idx], self.targets[idx].
indices = np.arange(len(self.dataset)), then, if self.shuffle, np.random.RandomState(self.seed).shuffle(indices) shuffles it IN PLACE (it returns None, not a shuffled copy).
for start in range(0, len(indices), self.batch_size): then batch_indices = indices[start : start + self.batch_size] naturally gives you a shorter final slice if the total doesn't divide evenly, no special-casing needed. Gather with np.stack([self.dataset[i][0] for i in batch_indices]) for the features (and similarly for targets), then yield the pair.
Click "Run Tests" to test your implementation