PyTorch random_split()返回大小错误的加载器

import torch
from torch.utils.data import DataLoader, Dataset

# 定义自定义的数据集类
class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        return self.data[index]

# 创建数据集
data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
dataset = CustomDataset(data)

# 计算划分的样本数量
dataset_size = len(dataset)
train_ratio = 0.8
train_size = int(dataset_size * train_ratio)
valid_size = dataset_size - train_size

# 划分数据集
train_dataset, valid_dataset = torch.utils.data.random_split(dataset, [train_size, valid_size])

# 创建加载器
batch_size = 2
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=True)

# 打印加载器的大小
print(f"Train loader size: {len(train_loader)}")
print(f"Valid loader size: {len(valid_loader)}")