diff --git a/pcdet/datasets/once/once_dataset.py b/pcdet/datasets/once/once_dataset.py
index 0ff079843..92d6b330c 100644
--- a/pcdet/datasets/once/once_dataset.py
+++ b/pcdet/datasets/once/once_dataset.py
@@ -12,6 +12,7 @@
 from ...utils import box_utils
 from .once_toolkits import Octopus
 
+
 class ONCEDataset(DatasetTemplate):
     def __init__(self, dataset_cfg, class_names, training=True, root_path=None, logger=None):
         """
@@ -388,8 +389,14 @@ def evaluation(self, det_annos, class_names, **kwargs):
 def create_once_infos(dataset_cfg, class_names, data_path, save_path, workers=4):
     dataset = ONCEDataset(dataset_cfg=dataset_cfg, class_names=class_names, root_path=data_path, training=False)
 
-    splits = ['train', 'val', 'test', 'raw_small', 'raw_medium', 'raw_large']
-    ignore = ['test']
+    image_sets = save_path / Path("ImageSets")
+    exist_sets = set(list(map(lambda p: p.stem, list(image_sets.glob("*.txt")))))
+
+    splits = set(['train', 'val', 'test', 'raw_small', 'raw_medium', 'raw_large'])
+    ignore = splits - exist_sets
+
+    splits = list(splits)
+    ignore = list(ignore)
 
     print('---------------Start to generate data infos---------------')
     for split in splits:
@@ -416,7 +423,7 @@ def create_once_infos(dataset_cfg, class_names, data_path, save_path, workers=4)
 
     parser = argparse.ArgumentParser(description='arg parser')
     parser.add_argument('--cfg_file', type=str, default=None, help='specify the config of dataset')
-    parser.add_argument('--func', type=str, default='create_waymo_infos', help='')
+    parser.add_argument('--func', type=str, default='create_once_infos', help='')
     parser.add_argument('--runs_on', type=str, default='server', help='')
     args = parser.parse_args()
 
@@ -441,4 +448,4 @@ def create_once_infos(dataset_cfg, class_names, data_path, save_path, workers=4)
             class_names=['Car', 'Bus', 'Truck', 'Pedestrian', 'Bicycle'],
             data_path=once_data_path,
             save_path=once_save_path
-        )
\ No newline at end of file
+        )
diff --git a/requirements.txt b/requirements.txt
index 40e445981..72dbed81c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,7 +4,7 @@ numba
 torch>=1.1
 tensorboardX
 easydict
-pyyaml
+pyyaml==5.4.1
 scikit-image
 tqdm
 torchvision