Updated load method in ModelBuilder

pymc_experimental/model_builder.py

@@ -158,7 +158,7 @@ def save(self, fname):
         self.idata.to_netcdf(file)
 
     @classmethod
-    def load(cls, fname):
+    def load(cls, self, fname):
         """
         Loads inference data for the model.
 
@@ -187,18 +187,15 @@ def load(cls, fname):
         filepath = Path(str(fname))
         data = az.from_netcdf(filepath)
         idata = data
-        # Since there is an issue with attrs getting saved in netcdf format which will be fixed in future the following part of code is commented
-        # Link of issue -> https://github.com/arviz-devs/arviz/issues/2109
-        # if model.idata.attrs is not None:
-        #     if model.idata.attrs['id'] == self.idata.attrs['id']:
-        #         self = model
-        #         self.idata = data
-        #         return self
-        #     else:
-        #         raise ValueError(
-        #             f"The route '{file}' does not contain an inference data of the same model '{self.__name__}'"
-        #         )
-        return idata
+        if idata.attrs is not None:
+            if self.id() == idata.attrs["id"]:
+                self = cls(idata.attrs["sample_config"], idata.attrs["model_config"])
+                self.idata = idata
+            else:
+                raise ValueError(
+                    f"The route '{file}' does not contain an inference data of the same model '{self.__name__}'"
+                )
+        return self
 
     def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None):
         """
@@ -238,8 +235,8 @@ def fit(self, data: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None
         self.idata.attrs["id"] = self.id()
         self.idata.attrs["model_type"] = self._model_type
         self.idata.attrs["version"] = self.version
-        self.idata.attrs["sample_conifg"] = self.sample_config
-        self.idata.attrs["model_config"] = self.model_config
+        self.idata.attrs["sample_config"] = tuple(self.sample_config)
+        self.idata.attrs["model_config"] = tuple(self.model_config)
         return self.idata
 
     def predict(