fix silly error in masked kl div loss, thanks to @taynoel84

lucidrains · Mar 12, 2023 · 2bb47d4 · 2bb47d4
1 parent 82fa3d0
commit 2bb47d4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/palm_rlhf_pytorch/ppo.py b/palm_rlhf_pytorch/ppo.py
@@ -265,7 +265,7 @@ def masked_kl_div(prob1, prob2, mask = None):
     """
     need to account for variable sequence lengths, therefore not using the built-in functional version
     """
-    kl_divs = (prob1 * (log(prob2) - log(prob1))).sum(dim = -1)
+    kl_divs = (prob1 * (log(prob1) - log(prob2))).sum(dim = -1)
 
     if not exists(mask):
         return kl_divs.mean()

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'PaLM-rlhf-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.66',
+  version = '0.0.67',
   license='MIT',
   description = 'PaLM + Reinforcement Learning with Human Feedback - Pytorch',
   author = 'Phil Wang',