def METHOD_NAME(self):
    train_ds = (
        RandomIterableDataset(32, count=max_steps + 100)
        if use_infinite_dataset
        else RandomDataset(32, length=data_samples_train)
    )
    return DataLoader(train_ds)