train_test_split.py

from torch.utils.data import random_split

def split_train_test(dataset, train_ratio=0.9):
    train_size = int(train_ratio * len(dataset))
    test_size = len(dataset) - train_size
    train_dataset, test_dataset = random_split(dataset, [train_size, test_size])
    return train_dataset, test_dataset