diff --git a/miles/ray/rollout_data_source.py b/miles/ray/rollout_data_source.py
index 649246ef469..7c47cdd8910 100644
--- a/miles/ray/rollout_data_source.py
+++ b/miles/ray/rollout_data_source.py
@@ -45,9 +45,7 @@ def __init__(self, args):
             self.dataset = None
 
     def get_samples(self, num_samples):
-        samples = []
-
-        # TODO unify the two branches
+        # TODO further improve code
         if self.dataset is not None:
             if self.sample_offset + num_samples <= len(self.dataset):
                 prompt_samples = self.dataset.samples[self.sample_offset : self.sample_offset + num_samples]
@@ -60,25 +58,18 @@ def get_samples(self, num_samples):
                     self.dataset.shuffle(self.epoch_id)
                 prompt_samples += self.dataset.samples[:num_samples]
                 self.sample_offset = num_samples
-            for prompt_sample in prompt_samples:
-                group = []
-                for _ in range(self.args.n_samples_per_prompt):
-                    sample = copy.deepcopy(prompt_sample)
-                    sample.index = self.sample_index
-                    self.sample_index += 1
-                    group.append(sample)
-                samples.append(group)
         else:
-            for _ in range(num_samples):
-                group = []
-                for _ in range(self.args.n_samples_per_prompt):
-                    sample = Sample(
-                        index=self.sample_index,
-                    )
-                    self.sample_index += 1
-                    group.append(sample)
-                samples.append(group)
+            prompt_samples = [Sample() for _ in range(num_samples)]
 
+        samples = []
+        for prompt_sample in prompt_samples:
+            group = []
+            for _ in range(self.args.n_samples_per_prompt):
+                sample = copy.deepcopy(prompt_sample)
+                sample.index = self.sample_index
+                self.sample_index += 1
+                group.append(sample)
+            samples.append(group)
         return samples
 
     def add_samples(self, samples: list[list[Sample]]):