正在尝试访问pytorch中mnist数据集的子集[每个类的样本相等]

import torch
from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为张量
    transforms.Normalize((0.5,), (0.5,))  # 标准化图像数据
])

train_dataset = datasets.MNIST('path_to_save_data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('path_to_save_data', train=False, download=True, transform=transform)

# 获取每个类的样本数量
class_counts = [0] * 10
for _, label in train_dataset:
    class_counts[label] += 1

# 设置每个类的子集样本数量
subset_size = min(class_counts)
subset_indices = []
for class_index in range(10):
    indices = [i for i, (_, label) in enumerate(train_dataset) if label == class_index]
    subset_indices.extend(indices[:subset_size])

# 创建子集数据集
subset_dataset = torch.utils.data.Subset(train_dataset, subset_indices)