from torch.utils.data import ConcatDataset
from tagging.data_loaders.ytm import YTMDataset


def get_dataset(dataset, split, input_length_s=30.0, num_iteration=1):
    if dataset == "ytm":
        return YTMDataset(
            split=split,
            input_length_s=input_length_s,
            num_iteration=num_iteration,
        )

    else:
        print("%s dataset is not supported yet" % dataset)