cleanup

lucidrains · Nov 18, 2024 · 5822175 · 5822175
1 parent 6cccff4
commit 5822175
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/pi_zero_pytorch/pi_zero.py b/pi_zero_pytorch/pi_zero.py
@@ -672,16 +672,14 @@ def forward_with_reward_cfg(
             return_actions_flow = True,
         )
 
-        maybe_reward_out = self.forward(
+        action_flow_with_reward, with_reward_cache_kv = self.forward(
             *args,
             reward_tokens = reward_tokens,
             cached_state_keys_values = with_reward_cache,
             **forward_kwargs,
             **kwargs
         )
 
-        action_flow_with_reward, with_reward_cache_kv = maybe_reward_out
-
         if not exists(reward_tokens) or cond_scale == 0.:
             return action_flow_with_reward, (with_reward_cache_kv, None)