diff --git a/examples/04_eval_finetuned_on_robot.py b/examples/04_eval_finetuned_on_robot.py
index ddd8e5d9..32cca713 100644
--- 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,