How to load from two datasets with different batch size

  I have two datasets, suppose they are A and B.  And I want to create one dataloader that can read from both of them with 512 samples from A and only 1 sample from B. And dataset A is the main one, the loop should end when A finishes its iteration. 
 Currently, my solution is below but it's time-consuming. Could any help me about this😣
dataset_A = lmdbDataset(*args)
dataset_B = lmdbDataset(*args
dataloader_A =, batch_size=512,shuffle=True)
dataloader_B =, batch_size=1,shuffle=True)
for batch_A in dataloader(dataset_A):
      batch_B = next(iter(dataloader_B)
      total_batch =[batch_A, batch_B], dim=0)
      out = model(total_batch)


first - your dataset is a part of your dataloader -
you do not need to mention ‘for batch_A in dataloader(dataset_A)’, but you can: for batchA in dataloader_A

fix this, and show what is your error flag and we can go from there