Hi everyone,
In my current setup, I would like to change the dataloader during a training epoch:

This is what I would like to achieve:
step 1.Train on dataset 1 for n batches
step 2.Train on dataset 2 for n batches
step 3.Go to step 1

I foundthis solution on the old forum but this only switches the dataset after each epoch.

Here is my current attempt at switching it every n batches:

classSimpleModule(pl.LightningModule):def__init__(self):super().__init__()self.model= ...self.batch_size= ...self.change_every_n_batch=20deftrain_dataloader(self):self.current_dataset= (self.global_step//self.change_every_n_batch)%2ifself.current_dataset==0:dataset=Dataset1()elifself.current_dataset==1:dataset=Dataset2()dataloader=DataLoader(dataset,batch_size=self.batch_size)returndataloaderdefon_train_batch_end(self,outputs,batch,batch_idx):new_dataset= (self.global_step//self.change_every_n_batch)%2ifnew_dataset!=self.current_dataset:self.trainer.reset_train_dataloader(self)

train_dataloader() is called as expected every 20 batches byon_train_batch_end() but the returned dataloader does not seem to be used during the training loop.

Any idea what could be going wrong? Or do you have a solution for what I want to achieve?

Thanks!