Are Lightning DataModules compatible with Modin?

[EvanZ]

I'm trying to use Modin as a drop-in replacement for Pandas. I have no problem creating a Dataset class based on this. But when I test the DataModule it just runs forever without returning anything:

class TestDataModule(LightningDataModule):
    def __init__(self, path:str=None, train_fraction:float=0.9, batch_size:int=32, workers:int=0) -> None:
        super().__init__()
        self._path = path
        self._train_fraction = train_fraction
        self._batch_size = batch_size
        self._workers = workers
        self.train_dataset = None
        self.val_dataset = None
        self.test_dataset = None
        
    def setup(self, stage: Optional[str]) -> None:
        if stage=='fit':
            dataset =TestDataset(self._path)
            self.train_dataset, self.val_dataset = train_test_split(dataset, train_size=self._train_fraction)
        if stage=='test':
            self.test_dataset = TestDataset(self._path)
            
    def train_dataloader(self) -> Union[DataLoader, List[DataLoader], Dict[str, DataLoader]]:
        return DataLoader(self.train_dataset, batch_size=self._batch_size, shuffle=True, pin_memory=True, \
            num_workers=self._workers)

    def val_dataloader(self) -> Union[DataLoader, List[DataLoader], Dict[str, DataLoader]]:
        return DataLoader(self.val_dataset, batch_size=self._batch_size, shuffle=True, pin_memory=True, \
            num_workers=self._workers)
    
    def test_dataloader(self) -> Union[DataLoader, List[DataLoader], Dict[str, DataLoader]]:
        return DataLoader(self.test_dataset, batch_size=self._batch_size, shuffle=False, pin_memory=True, \
            num_workers=self._workers)

if __name__=="__main__":
    parser =ArgumentParser()
    parser.add_argument("--path", help="path/to/files", type=str, default=None)
    args = parser.parse_args()

    datamodule = TestDataModule(path=args.path, batch_size=8)
    datamodule.setup(stage='fit')
    train_dataloader = datamodule.train_dataloader()
    print(type(train_dataloader))

Has anyone else had success using Modin in this way?