fix feature correlation model output(#5228)

Signed-off-by: Yu Wu <yolandawu131@gmail.com>
FederatedAI · Nov 22, 2023 · 5b14d69 · 5b14d69
1 parent 4da16d0
commit 5b14d69
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/python/fate/ml/statistics/pearson_correlation.py b/python/fate/ml/statistics/pearson_correlation.py
@@ -116,9 +116,10 @@ def fix_vif(remainds_vif, remainds_indexes, size):
         return vif"""
 
     def get_model(self):
-        output_model = {"data": {"local_corr": self.local_corr.to_dict(),
-                                 "remote_corr": self.remote_corr.to_dict(),
+        output_model = {"data": {"local_corr": self.local_corr.to_dict() if self.local_corr is not None else None,
+                                 "remote_corr": self.remote_corr.to_dict() if self.remote_corr is not None else None,
                                  "vif": self.vif.to_dict() if self.vif is not None else None},
                         "meta": {"model_type": "feature_correlation",
-                                 "column_anonymous_map": dict(zip(self.select_cols, self.select_anonymous_cols))}}
+                                 "column_anonymous_map": dict(zip(self.select_cols, self.select_anonymous_cols))
+                                 if self.select_anonymous_cols else None}}
         return output_model