如何在pytorch中对数据集进行排序

import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, data):
        self.data = data

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

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

def sort_dataset(dataset, sort_key):
    sorted_dataset = sorted(dataset, key=lambda x: x[sort_key])
    return sorted_dataset

data = [(1, 'A'), (3, 'C'), (2, 'B')]
dataset = CustomDataset(data)
sorted_dataset = sort_dataset(dataset, sort_key=0)

dataloader = DataLoader(sorted_dataset, batch_size=32, shuffle=True)