import pathlib
import numpy as np
import torch
from torch.utils.data import Dataset as TorchDataset


class ChessDataset(TorchDataset):
    def __init__(self, path: pathlib.Path):
        self.data = np.load(path)

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

    def __getitem__(self, index):
        return torch.tensor(self.data[index], dtype=torch.long)