model attributes correclty accessed

Author: Julien Roussel

qolmat/imputations/diffusions/ddpms.py

@@ -3,7 +3,7 @@
 import logging
 import time
 from datetime import timedelta
-from typing import Callable, Dict, List, Tuple, Union
+from typing import Any, Callable, Dict, List, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -132,7 +132,7 @@ def __init__(
         seed_torch = self.random_state.randint(2**31 - 1)
         torch.manual_seed(seed_torch)
 
-    def __getstate__(self) -> str:
+    def __getstate__(self) -> dict[str, Any]:
         """Hashing method used in sklearn check tests.
 
         Returns

qolmat/imputations/imputers_pytorch.py

@@ -876,7 +876,8 @@ def get_summary_training(self) -> Dict:
             Summary of the training
 
         """
-        return self.model.summary
+        model = self.get_model()
+        return model.summary
 
     def get_summary_architecture(self) -> Dict:
         """Get the summary of the architecture.
@@ -887,7 +888,8 @@ def get_summary_architecture(self) -> Dict:
             Summary of the architecture
 
         """
+        model = self.get_model()
         return {
-            "number_parameters": self.model.num_params,
-            "epsilon_model": self.model._eps_model,
+            "number_parameters": model.num_params,
+            "epsilon_model": model._eps_model,
         }