diff --git a/dataset/lm_dataset.py b/dataset/lm_dataset.py
index 1fa0397..c5fc48e 100644
--- a/dataset/lm_dataset.py
+++ b/dataset/lm_dataset.py
@@ -159,7 +159,7 @@ class DPODataset(Dataset):
                     if input_ids[end:end + len(self.eos_id)] == self.eos_id:
                         break
                     end += 1
-                for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):
+                for j in range(start, min(end + len(self.eos_id) + 1, self.max_length)):
                     loss_mask[j] = 1
                 i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
             else:
@@ -202,4 +202,4 @@ class RLAIFDataset(Dataset):
         }
 
 if __name__ == "__main__":
-    pass
\ No newline at end of file
+    pass