From 62cdaae25152de3a76c4b1fbf36b50b7c4b64440 Mon Sep 17 00:00:00 2001 From: chenyangkang Date: Sat, 26 Oct 2024 14:58:36 -0500 Subject: [PATCH] fix --- stemflow/model/AdaSTEM.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/stemflow/model/AdaSTEM.py b/stemflow/model/AdaSTEM.py index 43a1ea4..89e244d 100644 --- a/stemflow/model/AdaSTEM.py +++ b/stemflow/model/AdaSTEM.py @@ -556,9 +556,9 @@ def mp_train(ensemble, self=self, data=data): # iterate through if self.lazy_loading: - model_dict = LazyLoadingEnsembleDict(self.lazy_loading_dir) + self.model_dict = LazyLoadingEnsembleDict(self.lazy_loading_dir) else: - model_dict = {} + self.model_dict = {} stixel_specific_x_names = {} @@ -570,14 +570,13 @@ def mp_train(ensemble, self=self, data=data): name = feature_tuple[0] model = feature_tuple[1] x_names = feature_tuple[2] - model_dict[f"{name}_model"] = model + self.model_dict[f"{name}_model"] = model stixel_specific_x_names[name] = x_names # dump here if lazy_loading_ensemble = True if self.lazy_loading: - model_dict.dump_ensemble(ensemble_id) + self.model_dict.dump_ensemble(ensemble_id) - self.model_dict = model_dict self.stixel_specific_x_names = stixel_specific_x_names return self