Pytorch: How to make sure that all labels are present in each batch

Question

How to make sure that each batch will have samples with all the labels? For example, consider sentiment analysis problem with labels positive and negative.

tokens = tokenizer.batch_encode_plus(text.tolist(),max_length = max_seq_len,pad_to_max_length=True,truncation=True, return_token_type_ids=False)    
seq = torch.tensor(tokens['input_ids'])
mask = torch.tensor(tokens['attention_mask'])
y = torch.tensor(labels.tolist())    
data = TensorDataset(seq, mask,y)
data_sampler = RandomSampler(data)
data_dataloader = DataLoader(data, sampler=data_sampler, batch_size=batch_size)

I want to have batches like

Batch-1 ['positive','positive','positive','negative']
Batch-2 ['negative','negative','positive','negative']

where every batch contains all the labels.

Answer 1

There is this option in PyTorch about stratified sampling.

But if this does not satisfy your needs, my suggestion will be to either do it with scikit-learn adapting PyTorch code, or to read scikit-learn code and adapt it to PyTorch.