如何创建具有多个标签和掩码的自定义Pytorch数据集？

import torch
from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self, data, labels, masks):
        self.data = data
        self.labels = labels
        self.masks = masks

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

    def __getitem__(self, index):
        sample = {
            'data': self.data[index],
            'labels': self.labels[index],
            'masks': self.masks[index]
        }
        return sample

data = [...]  # 数据
labels = [...]  # 标签
masks = [...]  # 掩码

dataset = CustomDataset(data, labels, masks)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)