Method __init__

data/checkerboarddataset.py:236–244

(self, num_samples=5000, batch_size=1000, num_classes = 4, num_support = 200)

class CheckerboardDataModule(pl.LightningDataModule):
    def __init__(self, num_samples=5000, batch_size=1000, num_classes = 4, num_support = 200):
        super().__init__()
        self.num_samples = num_samples
        self.batch_size=batch_size
        self.num_support = num_support
        self.num_classes = num_classes

        # mean and std distance between clusters given the number of points
        self.mean_dist, self.std_dist = calc_avg_distances(num_support, unit="deg")

    def setup(self, stage: str):
        self.train_ds = TensorDataset(*get_data(N_samples = self.num_samples,