fix order of supply rng and partial in example

octo-models · May 20, 2024 · 86ff03a · 86ff03a
1 parent 4c731c5
commit 86ff03a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/04_eval_finetuned_on_robot.py b/examples/04_eval_finetuned_on_robot.py
@@ -129,8 +129,8 @@ def sample_actions(
         # remove batch dim
         return actions[0]
 
-    policy_fn = partial(
-        supply_rng(
+    policy_fn = supply_rng(
+        partial(
             sample_actions,
             model,
             argmax=FLAGS.deterministic,