Minor refactors based on review

GFNOrg · Dec 6, 2024 · 5ce1fdc · 5ce1fdc
1 parent 21ce0c2
commit 5ce1fdc
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/src/gfn/samplers.py b/src/gfn/samplers.py
@@ -247,10 +247,10 @@ def sample_trajectories(
 
         trajectories_states = stack_states(trajectories_states)
         trajectories_actions = env.Actions.stack(trajectories_actions)[
-            1:, ...  # Drop dummy action
+            1:  # Drop dummy action
         ]
         trajectories_logprobs = (
-            torch.stack(trajectories_logprobs, dim=0)[1:, ...]  # Drop dummy logprob
+            torch.stack(trajectories_logprobs, dim=0)[1:]  # Drop dummy logprob
             if save_logprobs
             else None
         )
@@ -546,7 +546,7 @@ def sample_trajectories(
             n = trajectories.n_trajectories
 
         search_indices = torch.arange(n, device=trajectories.states.device)
-        for it in range(n_local_search_loops - 1):
+        for it in range(1, n_local_search_loops):  # 0-th loop is the initial sampling
             # Search phase
             ls_trajectories, is_updated = self.local_search(
                 env,
@@ -562,7 +562,7 @@ def sample_trajectories(
             trajectories.extend(ls_trajectories)
 
             last_indices = torch.arange(
-                n * (it + 1), n * (it + 2), device=trajectories.states.device
+                n * it, n * (it + 1), device=trajectories.states.device
             )
             search_indices[is_updated] = last_indices[is_updated]
 

diff --git a/tutorials/examples/train_hypergrid_simple_ls.py b/tutorials/examples/train_hypergrid_simple_ls.py
@@ -38,7 +38,7 @@ def main(args):
     gflownet = TBGFlowNet(pf=pf_estimator, pb=pb_estimator, logZ=0.0)
 
     # Feed pf to the sampler.
-    sampler = LocalSearchSampler(estimator=pf_estimator, pb_estimator=pb_estimator)
+    sampler = LocalSearchSampler(pf_estimator=pf_estimator, pb_estimator=pb_estimator)
 
     # Move the gflownet to the GPU.
     gflownet = gflownet.to(device_str)