diff --git a/code/gpt_loader/load_data.py b/code/gpt_loader/load_data.py
index db508be66733a91b3e99139b35955cfdc5845542..b5e78b8de001bdd417d9c2e2533f2e06bbe4f013 100644
--- a/code/gpt_loader/load_data.py
+++ b/code/gpt_loader/load_data.py
@@ -480,7 +480,7 @@ class GptDataset_KBERT(Dataset):
             soft_position_x += list(range(soft_loc, soft_loc + len(utternace_encoded) + 1))
 
             # add the aug, if it is the right place
-            if len(dq) != 0 and dq[0][0] == i:
+            while len(dq) != 0 and dq[0][0] == i:
                 comet_output = dq.popleft()[1]
                 comet_encoded = self.tokenizer.encode(text_standardize(comet_output))
                 x += [self.augment] + comet_encoded