Use two datasets and distinguish during training

Hi, I am aware of multiple loaders but I still can’t figure it out for what I need to do. I have a dataset of graphs (like mutag) and then I computed their line graphs, so in the end I have two datasets and they are divided in batches. In the training I am running the same procedure for both datasets but I need to run different methods in the forward step, depending on the origin. I don’t get how the data is managed, is it possible to distinguish between the two datasets in that step? If so how can I do it? Do I need to use the method I linked? I hope that what I wrote is understandable. Thank you.