How to load from two datasets with different batch size

vision
1 
  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 = torch.utils.data.Dataloader(dataset_A, batch_size=512,shuffle=True)
dataloader_B = torch.utils.data.Dataloader(dataset_B, batch_size=1,shuffle=True)
...
for batch_A in dataloader(dataset_A):
      batch_B = next(iter(dataloader_B)
      total_batch = torch.cat([batch_A, batch_B], dim=0)
      out = model(total_batch)
2

Hi-

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