diff --git a/tutorials/examples/train_hypergrid_gafn.py b/tutorials/examples/train_hypergrid_gafn.py
index 936cdc7b..7ad43cc6 100644
--- a/tutorials/examples/train_hypergrid_gafn.py
+++ b/tutorials/examples/train_hypergrid_gafn.py
@@ -162,7 +162,10 @@ def flow_parameters(self) -> list[torch.Tensor]:
         )
 
     def get_scores(
-        self, trajectories: Trajectories, recalculate_all_logprobs: bool = True
+        self,
+        trajectories: Trajectories,
+        recalculate_all_logprobs: bool = True,
+        log_rewards: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Computes Trajectory Balance scores with intrinsic rewards for a batch of
         trajectories.
@@ -170,6 +173,8 @@ def get_scores(
         Args:
             trajectories: The Trajectories object to evaluate.
             recalculate_all_logprobs: Whether to re-evaluate all logprobs.
+            log_rewards: Optional override for the trajectories' log rewards. If None,
+                defaults to `trajectories.log_rewards`.
 
         Returns:
             A tensor of shape (batch_size,) containing the scores for each trajectory.
@@ -177,7 +182,8 @@ def get_scores(
         log_pf_trajectories, log_pb_trajectories = self.get_pfs_and_pbs(
             trajectories, recalculate_all_logprobs=recalculate_all_logprobs
         )
-        log_rewards = trajectories.log_rewards
+        if log_rewards is None:
+            log_rewards = trajectories.log_rewards
         assert log_rewards is not None
         if self.use_edge_ri:
             # Use the edge-based intrinsic rewards.