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


features = torch.arange(24, dtype=torch.float32).reshape(6, 4)
labels = torch.arange(6)
dataset = TensorDataset(features, labels)