Merge pull request #29 from Control-derek/derek

Fixed deep copy, shallow copy error and label mask error.
lucidrains · Apr 11, 2024 · d4755a2 · d4755a2
2 parents 2db4fed + 2c3ec61
commit d4755a2
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/self_rewarding_lm_pytorch/self_rewarding_lm_pytorch.py b/self_rewarding_lm_pytorch/self_rewarding_lm_pytorch.py
@@ -315,9 +315,9 @@ def get_cross_entropy_loss(
         else:
             prompt_mask = prompt_len_or_mask
 
-        seq, labels = seq[:, :-1], seq[:, 1:]
+        seq, labels = seq[:, :-1].clone(), seq[:, 1:].clone()
 
-        labels.masked_fill_(prompt_mask[:, 1:], self.ignore_index)
+        labels.masked_fill_(~prompt_mask[:, 1:], self.ignore_index)
 
         logits = self.model(seq)
 
@@ -822,6 +822,7 @@ def __init__(
                     eval_filter_fn = config.eval_filter_fn,
                     eval_filter_kwargs = config.eval_filter_kwargs,
                     accelerator = self.accelerator,
+                    pad_id = pad_id,
                     **config.reward_generator_kwargs
                 )
 
@@ -862,6 +863,7 @@ def __init__(
                     eval_temperature = config.eval_temperature,
                     eval_filter_fn = config.eval_filter_fn,
                     eval_filter_kwargs = config.eval_filter_kwargs,
+                    pad_id = pad_id,
                     **config.reward_generator_kwargs
                 )