fix masking logic when using palm as encoder

lucidrains · Jan 8, 2023 · 6e1baf6 · 6e1baf6
1 parent 4cac0e6
commit 6e1baf6
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/palm_rlhf_pytorch/palm_rlhf_pytorch.py b/palm_rlhf_pytorch/palm_rlhf_pytorch.py
@@ -477,7 +477,11 @@ def forward(
         # mask if encoder
         # treat any token ids that are negative as tokens to mask out - only needed if not autoregressive
 
-        mask = (x < 0) if not self.causal else None
+        if not self.causal:
+            mask = x >= 0
+            x = x.masked_fill(~mask, 0)
+        else:
+            mask = None
 
         # get token embedding
 

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