Custom Dataset to sample imbalanced dataset

I have a dataset with data that falls into one of three labels/classes : A, B, C

class A has 3000 data points
class B has 2000 data points
class C has 1000 data points.

In every epoch of my training I want my Dataset to pick all 1000 data points of C and randomly select 1000 data points of B and 1000 data points of A .

As it stands in the code below, the Dataloader will parse through all 6000 data points. How do I modify the Dataset below to ‘balance’ the dataset so that I use a 1000 points from each class at every epoch?

A = np.random.rand(3000)
B = np.random.rand(2000)
C = np.random.rand(1000)

data_array = np.concatenate([A,B,C])
labels = np.concatenate([np.zeros_like(A), np.ones_like(B), 2*np.ones_like(C)])

class Imbalanced_Dataset(Dataset):
    def __init__(self, data_array, labels):
        self.data_array = data_array
        self.labels = labels
    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):

        data_point = self.data_list[index]
        label = self.label[index]

        return data_point,label