For instance dataset.train_test_split(...) returns a HF Dataset, which then breaks serialization, deserialization logic.
Not sure if there's a better solution but the quick fix would be to wrap the methods (train_test_split, filter, select, sort, etc.) with Dataset.from_hf_dataset(...)