Merge pull request #13 from Dyke-F/main

Update self_rewarding_lm_pytorch.py
lucidrains · Jan 30, 2024 · e09fccf · e09fccf
2 parents 7ecb8a7 + 710d611
commit e09fccf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/self_rewarding_lm_pytorch/self_rewarding_lm_pytorch.py b/self_rewarding_lm_pytorch/self_rewarding_lm_pytorch.py
@@ -331,7 +331,7 @@ def forward(self):
 
                                 loss = self.get_cross_entropy_loss(seq, prompt_len_or_mask)
 
-                                total_valid_loss = loss.item() * batch
+                                total_valid_loss += loss.item() * batch
                                 total_batches += batch
 
                         valid_loss = total_valid_loss / total_batches