Fixes (cherry-pick)

alexhernandezgarcia · Nov 28, 2023 · a777e99 · a777e99
1 parent 20285ff
commit a777e99
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/gflownet/gflownet.py b/gflownet/gflownet.py
@@ -808,10 +808,10 @@ def forwardlooking_loss(self, it, batch):
             - logprobs_b
             + energies_transitions
         ).pow(2)
-        loss = per_node_loss.mean()
-        loss_terminating = per_node_loss[done].mean()
-        loss_intermediate = per_node_loss[~done].mean()
-        return loss, term_loss, nonterm_loss
+        loss = loss_all.mean()
+        loss_terminating = loss_all[done].mean()
+        loss_intermediate = loss_all[~done].mean()
+        return loss, loss_terminating, loss_intermediate
 
     @torch.no_grad()
     def estimate_logprobs_data(