make sure critic is inferenced without grad and eval is turn on

lucidrains · Oct 28, 2022 · c1b1236 · c1b1236
1 parent b9de5d9
commit c1b1236
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/phenaki_pytorch/phenaki_pytorch.py b/phenaki_pytorch/phenaki_pytorch.py
@@ -967,6 +967,9 @@ def __init__(
 
         # sampling
 
+        if exists(critic):
+            critic = critic.eval()
+
         self.critic = critic
         self.steps = steps
         self.sample_temperature = sample_temperature
@@ -1060,7 +1063,8 @@ def sample(
 
             if not is_last_step:
                 if exists(self.critic):
-                    scores = self.critic(video_token_ids)
+                    with torch.no_grad():
+                        scores = self.critic(video_token_ids)
 
                     noise = K * (uniform(scores.shape, device) - 0.5) * (steps_til_x0 / self.steps)
                     scores = scores + noise

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