I am using a model where I want to initialize some random state in my Dataset that decides how many context points to add to the current instance (from this paper: https://arxiv.org/pdf/1807.01613.pdf)
I can’t figure out how to call a function on each batch. The only requirements I have is that

It can be seen by the Dataset object so I can use it in its __getitem__ method.
It will be called once for every batch.

Is there a way to accomplish this?