Need tfds to be pickable

Hi! I am struggling to make some torch distributed code to work with a tfds dataset.

I seem to have reduced the issue to the following few lines of code

Basically, if I could manage to get the tfds dataset (above) pickable I believe the torch distributed code would run. Is there a way to do that?

Here is the dataset btw: Directory listing for /tensorflow_datasets_small/cms_pf_ttbar/1.6.0/

Thanks for any help!!

solved it by defining this part to avoid pickling the dataset_info part

        tmp = self.ds.dataset_info
        from types import SimpleNamespace

        self.ds.dataset_info = SimpleNamespace()
        self.ds.dataset_info.name = tmp.name
        self.ds.dataset_info.features = tmp.features

will leave it here in case it helps anyone!