# Source code for archai.datasets.limit_dataset

    # Copyright (c) Microsoft Corporation.
    # Licensed under the MIT license.

    from typing import List, Tuple, Union, Optional

    import torch
    from torch.utils.data import \
        SubsetRandomSampler, Sampler, Subset, ConcatDataset, Dataset, random_split

    class LimitDataset(Dataset):
        def __init__(self, dataset, n):
            self.dataset = dataset
            self.n = n
            if hasattr(dataset, 'targets'):
                self.targets = dataset.targets[:n]

        def __len__(self):
            return self.n

        def __getitem__(self, i):
            return self.dataset[i]

    DatasetLike = Optional[Union[Dataset, Subset, ConcatDataset, LimitDataset]]