Merge pull request #97 from ChEB-AI/fix/save_out_dim_to_checkpoint

Fix `input_dim` and `out_dim` Not getting saved in model checkpoint

Author: sfluegel05

chebai/models/base.py

@@ -45,6 +45,11 @@ def __init__(
         if exclude_hyperparameter_logging is None:
             exclude_hyperparameter_logging = tuple()
         self.criterion = criterion
+        assert out_dim is not None, "out_dim must be specified"
+        assert input_dim is not None, "input_dim must be specified"
+        self.out_dim = out_dim
+        self.input_dim = input_dim
+
         self.save_hyperparameters(
             ignore=[
                 "criterion",
@@ -55,10 +60,8 @@ def __init__(
             ]
         )
 
-        self.out_dim = out_dim
-        self.input_dim = input_dim
-        assert out_dim is not None, "out_dim must be specified"
-        assert input_dim is not None, "input_dim must be specified"
+        self.hparams["out_dim"] = out_dim
+        self.hparams["input_dim"] = input_dim
 
         if optimizer_kwargs:
             self.optimizer_kwargs = optimizer_kwargs