a minor refactoring

GFNOrg · Nov 15, 2024 · b506769 · b506769
1 parent 92301c4
commit b506769
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/gfn/gflownet/detailed_balance.py b/src/gfn/gflownet/detailed_balance.py
@@ -157,7 +157,7 @@ def get_scores(
                 valid_log_F_s_next = self.logF(valid_next_states).squeeze(-1)
 
         log_F_s_next = torch.zeros_like(log_pb_actions)
-        log_F_s_next[~valid_transitions_is_done] += valid_log_F_s_next
+        log_F_s_next[~valid_transitions_is_done] = valid_log_F_s_next
         assert transitions.log_rewards is not None
         valid_transitions_log_rewards = transitions.log_rewards[
             ~transitions.states.is_sink_state