Searchable parameter in PyTorch Lightning DataModule?

I am using Ray Tune to perform HP search with my PyTorch Lightning project. I have my dataloaders in a PTL DataModule. However when building Ray’s LightningConfigBuilder, I have to pass an instance of my data-module for the .fit_params. Therefore I can’t use a Ray search space for my dataloader’s batch_size parameter.

Is there a way around this? I know I could move my dataloaders inside of my PTL LightningModule, but I would rather keep them in my PTL DataModule to keep the code modular/clean.


I had the same use case and was advised to move to their new version and use TorchTrainer.

Installation: Here’s the installation guide (install the Daily Releases (Nightlies) wheel). It’s from 2 weeks ago, so it might just be the stable version.

API: Lightning API guides

It’s a lot more flexible and boils down to passing your complete training function to the TorchTrainer object.

@f2010126 is right! We recommend using a TorchTrainer and creating the pl.Trainer manually in your custom training loop – this way you have full control over how to initialize your data loaders before training.

You can pass this into a Tuner and tune the batch size, then initialize the dataloaders accordingly at the start of your training function.

Take a look at the new Pytorch Lightning + Ray Train user guide: