Isn’t this the same issue as described in your previous post?
In your code snippet it seems you are still using the PyG DataLoader:
from torch_geometric import loader
...
train_loader = loader.DataLoader(train_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_my_dataset)