random_split 之后如何更改数据集中的目标？

Question

我有一个用于训练和测试的数据集，如下所示，

dataset['train'], dataset['test'] = torch.utils.data.random_split(dataset_all, [num_train,num_test],
                                                             generator=torch.Generator().manual_seed(random_seed))

有什么好的方法可以通过提供索引来更改目标并检索某些数据集吗？ 现在，我使用列表这种方式只是为了获取标签==0的数据集，

dataloader['train'] = torch.utils.data.DataLoader(dataset['train'], batch_size=len(dataset['train']), num_workers=4)
inputs, labels = next(iter(dataloader['train']))
x_train = inputs[np.where(labels==0)]
y_train = labels[np.where(labels==0)]
data_train = My_Dataset(x_train, y_train, transform=None)

当数据集大小太大时，这种方式会花费大量时间和内存。