Here is an example of custom batch_sampler for your case
def chunk(indices, size):
    return torch.split(torch.tensor(indices), size)

class MyBatchSampler(Sampler):
    def __init__(self, a_indices, b_indices, batch_size): 
        self.a_indices = a_indices
        self.b_indices = b_indices
        self.batch_size = batch_size
    
    def __iter__(self):
        random.shuffle(self.a_indices)
        random.shuffle(self.b_indices)
        a_batches  = chunk(self.a_indices, self.batch_size)
        b_batches = chunk(self.b_indices, self.batch_size)
        all_batches = list(a_batches + b_batches)
        all_batches = [batch.tolist() for batch in all_batches]
        random.shuffle(all_batches)
        return iter(all_batches)
   

new_dataset = ConcatDataset((dataset_a, dataset_b))
a_len = dataset_a.__len__()
ab_len = a_len + dataset_b.__len__()
a_indices = list(range(a_len))
b_indices = list(range(a_len, b_len))

batch_sampler = MyBatchSampler(a_indices, b_indices, batch_size)

dl = DataLoader(new_dataset,  batch_sampler=batch_sampler)

to verify if each batch only contains elements from either A or B
for x in batch_sampler:
    print(x)