Using Lightning DataModules

[simon-forb]

I have the following minimal code example:

from lightning import LightningDataModule
from torch.utils.data import random_split
from torch_geometric.data.lightning import LightningDataset
from torch_geometric.datasets import TUDataset
import lightning as L

trainer = L.Trainer()
reddit_dataset = TUDataset("./data", "REDDIT-BINARY")
train_set, val_set, test_set = random_split(reddit_dataset, [.7, .1, .2])
dm = LightningDataset(train_set, val_set, test_set, batch_size=32)

trainer.fit(model, datamodule=dm)

Which gives me the following error:

TypeError: An invalid dataloader was passed to `Trainer.fit(train_dataloaders=...)`. Found LightningDataset(train_dataset=<torch.utils.data.dataset.Subset object at 0x773852bfe710>, val_dataset=<torch.utils.data.dataset.Subset object at 0x773852c35110>, test_dataset=<torch.utils.data.dataset.Subset object at 0x773852bffa50>, batch_size=32, num_workers=0, pin_memory=True, persistent_workers=False).

I would have expected this is how LightningDataset is supposed to be applied.

My code works when I create standard DataLoaders from PyG.

Version: pyg 2.5.2 py311_torch_2.2.0_cu121

[rusty1s]

Which PL version are you using?

[simon-forb]

This is a very recent environment, it uses

pytorch-lightning=2.2.2=pyhd8ed1ab_0

[brenerrr]

I am having the same issue.

torch: 2.1.1
lightning: 2.4.0
torch_geometric: 2.4.0

[akihironitta]

Mind updating it to torch_geometric==2.6.0 and retrying it? I cannot reproduce this with the PyTorch Lightning version.