black

GFNOrg · Apr 2, 2024 · 9ae95a5 · 9ae95a5
1 parent 89c72b5
commit 9ae95a5
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/src/gfn/gflownet/base.py b/src/gfn/gflownet/base.py
@@ -168,7 +168,10 @@ def get_pfs_and_pbs(
         if has_log_probs(trajectories) and not recalculate_all_logprobs:
             log_pf_trajectories = trajectories.log_probs
         else:
-            if trajectories.estimator_outputs is not None and not recalculate_all_logprobs:
+            if (
+                trajectories.estimator_outputs is not None
+                and not recalculate_all_logprobs
+            ):
                 estimator_outputs = trajectories.estimator_outputs[
                     ~trajectories.actions.is_dummy
                 ]

diff --git a/src/gfn/gflownet/trajectory_balance.py b/src/gfn/gflownet/trajectory_balance.py
@@ -42,7 +42,10 @@ def __init__(
         self.log_reward_clip_min = log_reward_clip_min
 
     def loss(
-        self, env: Env, trajectories: Trajectories, recalculate_all_logprobs: bool = False
+        self,
+        env: Env,
+        trajectories: Trajectories,
+        recalculate_all_logprobs: bool = False,
     ) -> TT[0, float]:
         """Trajectory balance loss.
 
@@ -83,7 +86,10 @@ def __init__(
         self.log_reward_clip_min = log_reward_clip_min
 
     def loss(
-        self, env: Env, trajectories: Trajectories, recalculate_all_logprobs: bool = False
+        self,
+        env: Env,
+        trajectories: Trajectories,
+        recalculate_all_logprobs: bool = False,
     ) -> TT[0, float]:
         """Log Partition Variance loss.