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


class ChurnDataset(Dataset):
    def __init__(self, feature_rows, label_rows):
        if len(feature_rows) != len(label_rows):
            raise ValueError("feature_rows and label_rows must have the same length")
        self.features = torch.as_tensor(feature_rows, dtype=torch.float32)
        self.labels = torch.as_tensor(label_rows, dtype=torch.long)