Reload train dataloader when using limit_train_batches

[mrpep]

Hi, I am using an IterableDataset that generates infinite samples and I want it to change its behaviour as the training progresses. To do so, I need to access information about the current training step from inside the Dataset.
I tried passing the trainer to the dataset as an argument, however, when using multiple workers the trainer is not updated.
Another idea was to define virtual epochs of N steps. To do so, I passed limit_train_batches = N as an argument to the Trainer. Also, when an epoch finishs, I need to update the dataset variables. To do it I set reload_dataloader_every_n_epochs = 1 and I wrote a train_dataloader method in my Lightning Module that returns a new dataloader with the updated variables. However, it seems that the train_dataloader method isn't called if I set the limit_train_batches variable. Any idea on how to solve this problem?

Thanks in advance.

[rohitgr7]

However, it seems that the train_dataloader method isn't called if I set the limit_train_batches variable. Any idea on how to solve this problem?

this should not happen. Just tried a simple example using BoringModel and it is recalling the train dataloader hook correctly even when using limit_train_batches.

[mrpep]

I finally could solve it. The problem was that I passed the train and validation dataloaders as arguments to trainer.fit(), so train_dataloader() wasn't called. I ended up implementing __len__() in the IterableDataset instead of using limit_train_batches