should be using 1 - probs for score

lucidrains · Jan 19, 2023 · db96595 · db96595
1 parent 00fcc9f
commit db96595
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/phenaki_pytorch/phenaki_pytorch.py b/phenaki_pytorch/phenaki_pytorch.py
@@ -663,7 +663,8 @@ def sample(
                     noise = noise_K * (uniform(scores.shape, device) - 0.5) * noise_multiplier
                     scores = scores + noise
                 else:
-                    scores = logits.gather(2, rearrange(pred_video_ids, '... -> ... 1'))
+                    probs = logits.softmax(dim = -1)
+                    scores = probs.gather(2, rearrange(pred_video_ids, '... -> ... 1'))
                     scores = 1 - rearrange(scores, '... 1 -> ...')
                     scores = torch.where(mask, scores, -1e4)
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'phenaki-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.70',
+  version = '0.0.71',
   license='MIT',
   description = 'Phenaki - Pytorch',
   author = 'Phil Wang',