WeightedRandomSampler might be helpful to you.
Here’s a discussion about it.

  
    
    
    Balanced Sampling between classes with torchvision DataLoader vision
  
  
    Ok, I found some other issues. I changed your code a bit. This should work: 
numDataPoints = 1000
data_dim = 5
bs = 100

data = torch.FloatTensor(numDataPoints, data_dim)
target = np.hstack((np.zeros(int(numDataPoints * 0.9), dtype=np.int32),
                    np.ones(int(numDataPoints * 0.1), dtype=np.int32)))

dataset_x = data.numpy()
dataset_y = target
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(dataset_x,
                        …