Adapt Decision Transformer collator to handle evaluation

Hi, I am referring to this tutorial notebook. I was trying to add an evaluation to the training but now I am encountering some difficulties. It seems that in the original code the Collator does also a custom sampling, but this sampling will always be done on the train_set since at initialization I will use that one, as shown in the image

My question is: Is there a way to use 2 different collator logics based on if the model is evaluating or training? Another option I think would be using the collator on the 2 datasets before training but I don’t know if it is possible.