Fix CI: accept log_rewards kwarg in TBGAFN.get_scores

tutorials/examples/train_hypergrid_gafn.py

@@ -162,22 +162,28 @@ 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.
 
         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.
         """
         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.