diff options
Diffstat (limited to 'data/custom_dataset_data_loader.py')
| -rwxr-xr-x | data/custom_dataset_data_loader.py | 11 |
1 files changed, 9 insertions, 2 deletions
diff --git a/data/custom_dataset_data_loader.py b/data/custom_dataset_data_loader.py index 0b98254..89ae554 100755 --- a/data/custom_dataset_data_loader.py +++ b/data/custom_dataset_data_loader.py @@ -4,8 +4,15 @@ from data.base_data_loader import BaseDataLoader def CreateDataset(opt): dataset = None - from data.aligned_dataset import AlignedDataset - dataset = AlignedDataset() + if opt.phase == 'recursive': + from data.recursive_dataset import RecursiveDataset + dataset = RecursiveDataset() + elif opt.phase == 'sequence': + from data.sequence_dataset import SequenceDataset + dataset = SequenceDataset() + else: + from data.aligned_dataset import AlignedDataset + dataset = AlignedDataset() print("dataset [%s] was created" % (dataset.name())) dataset.initialize(opt) |