recursive dataset stuff

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)