diff --git a/model/reward/instructor/rank_datasets.py b/model/reward/instructor/rank_datasets.py
index 2f2260c2..c2b7e58f 100644
--- a/model/reward/instructor/rank_datasets.py
+++ b/model/reward/instructor/rank_datasets.py
@@ -120,13 +120,14 @@ class HFSummary(Dataset):
         conf_threshold=-1,
         max_comparison_per_sample=3) -> None:
         super().__init__()
-        assert split in ('train', 'validation')
+        assert split in ('train', 'valid1', 'valid2', 'test')
         summaries = {}
         # using prompt as our index will allows us
         # to add additional generated prompt later
         self.index2summary = {}
         self.max_comparison_per_sample = max_comparison_per_sample
-        dataset = load_dataset('Tristan/summarize_from_feedback', 'comparisons')[split]
+        major_split = split if 'train' == split else 'validation'
+        dataset = load_dataset('Tristan/summarize_from_feedback', 'comparisons')[major_split]
         for data in dataset:
             if 'extra' in data and \
                 'confidence' in data['extra'] and \
@@ -135,6 +136,9 @@ class HFSummary(Dataset):
                 print('skipping {}'.format(data['info']['id']))
                 continue
 
+            if split != 'train' and split != data['split']:
+                continue
+
             if 'article' in data['info'] and \
                 data['info']['article'] is not None:
                 context = data['info']['article']
diff --git a/model/reward/instructor/tests/test_dataset.py b/model/reward/instructor/tests/test_dataset.py
index 7b432fd3..5765cd43 100644
--- a/model/reward/instructor/tests/test_dataset.py
+++ b/model/reward/instructor/tests/test_dataset.py
@@ -7,7 +7,7 @@ def test_hfsummary():
     
     tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
     collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
-    dataset = HFSummary()
+    dataset = HFSummary('train')
     print(len(dataset))
     dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8)
     for batch in dataloader: