AttributeError: 'DatasetDict' object has no attribute 'train_test_split'

[alan-ai]

or this can also works

squad = load_dataset('squad')['train'].train_test_split(train_size=800, test_size=200)