Fix config issue by implementing _get_config()

gflownet/policy/base.py

@@ -36,9 +36,7 @@ def __init__(
         float_precision : int or torch.dtype
             The floating point precision to be passed to torch tensors.
         """
-        # If config is None, instantiate an empty config (defaults will be used)
-        if config is None:
-            config = OmegaConf.create()
+        config = self._get_config(config)
         # Device and float precision
         self.device = set_device(device)
         self.float = set_float_precision(float_precision)
@@ -56,6 +54,26 @@ def __init__(
         # Instantiate the model
         self.model, self.is_model = self.make_model()
 
+    @staticmethod
+    def _get_config(config: Union[dict, DictConfig]) -> Union[dict, DictConfig]:
+        """
+        Returns a configuration dictionary, even if the input is None.
+
+        Parameters
+        ----------
+        config : dict or DictConfig
+            The configuration dictionary to set up the policy model. It may be None, in
+            which an empty config is created and the defaults will be used.
+
+        Returns
+        -------
+        config : dict or DictConfig
+            The configuration dictionary to set up the policy model.
+        """
+        if config is None:
+            config = OmegaConf.create()
+        return config
+
     def make_model(self) -> Tuple[Union[torch.Tensor, torch.nn.Module], bool]:
         """
         Instantiates the model of the policy.

gflownet/policy/cnn.py

@@ -7,6 +7,7 @@
 
 class CNNPolicy(Policy):
     def __init__(self, **kwargs):
+        config = self._get_config(kwargs["config"])
         # Shared weights, defaults to False
         self.shared_weights = config.get("shared_weights", False)
         # Reload checkpoint, defaults to False

gflownet/policy/mlp.py

@@ -6,6 +6,7 @@
 
 class MLPPolicy(Policy):
     def __init__(self, **kwargs):
+        config = self._get_config(kwargs["config"])
         # Shared weights, defaults to False
         self.shared_weights = config.get("shared_weights", False)
         # Reload checkpoint, defaults to False