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

features = torch.arange(60, dtype=torch.float32).reshape(20, 3)
labels = torch.arange(20)
dataset = TensorDataset(features, labels)