Fix: logflows of terminating is log(rewards)

alexhernandezgarcia · Nov 27, 2023 · f27a826 · f27a826
1 parent 5a135d9
commit f27a826
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/gflownet/gflownet.py b/gflownet/gflownet.py
@@ -731,7 +731,7 @@ def detailedbalance_loss(self, it, batch):
 
         # Get logflows
         logflow_states = self.state_flow(states_policy).squeeze()
-        logflow_states[done.eq(1)] = rewards
+        logflow_states[done.eq(1)] = torch.log(rewards)
         # TODO: Optimise by reusing logflow_states and batch.get_parent_indices
         logflow_parents = self.state_flow(parents_policy).squeeze()