conditional gfn

src/gfn/modules.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Optional, Any
 
 import torch
 import torch.nn as nn
@@ -109,7 +109,7 @@ def to_probability_distribution(
         self,
         states: States,
         module_output: TT["batch_shape", "output_dim", float],
-        **policy_kwargs: Optional[dict],
+        **policy_kwargs: Any,
     ) -> Distribution:
         """Transform the output of the module into a probability distribution.
 
@@ -240,13 +240,20 @@ def __init__(
         self.conditioning_module = conditioning_module
         self.final_module = final_module
 
-    def forward(
-        self, states: States, conditioning: torch.tensor
+    def _forward_trunk(
+        self, states: States, conditioning: torch.Tensor
     ) -> TT["batch_shape", "output_dim", float]:
         state_out = self.module(self.preprocessor(states))
         conditioning_out = self.conditioning_module(conditioning)
         out = self.final_module(torch.cat((state_out, conditioning_out), -1))
 
+        return out
+
+    def forward(
+        self, states: States, conditioning: torch.tensor
+    ) -> TT["batch_shape", "output_dim", float]:
+        out = self._forward_trunk(states, conditioning)
+
         if not self._output_dim_is_checked:
             self.check_output_dim(out)
             self._output_dim_is_checked = True
@@ -272,13 +279,24 @@ def __init__(
             is_backward=is_backward,
         )
 
+    def forward(
+        self, states: States, conditioning: torch.tensor
+    ) -> TT["batch_shape", "output_dim", float]:
+        out = self._forward_trunk(states, conditioning)
+
+        if not self._output_dim_is_checked:
+            self.check_output_dim(out)
+            self._output_dim_is_checked = True
+
+        return out
+
     def expected_output_dim(self) -> int:
         return 1
 
     def to_probability_distribution(
         self,
         states: States,
         module_output: TT["batch_shape", "output_dim", float],
-        **policy_kwargs: Optional[dict],
+        **policy_kwargs: Any,
     ) -> Distribution:
         raise NotImplementedError