diff --git a/python/raydp/spark/ray_cluster.py b/python/raydp/spark/ray_cluster.py
index 74b4b536..5977b720 100644
--- a/python/raydp/spark/ray_cluster.py
+++ b/python/raydp/spark/ray_cluster.py
@@ -125,9 +125,10 @@ def _prepare_spark_configs(self):
 
         raydp_agent_path = os.path.abspath(os.path.join(os.path.abspath(__file__),
                                                         "../../jars/raydp-agent*.jar"))
-        print(raydp_agent_path)
-        raydp_agent_jar = glob.glob(raydp_agent_path)[0]
-        self._configs[SPARK_JAVAAGENT] = raydp_agent_jar
+        print(os.listdir(raydp_cp))
+        raydp_agent_jars = glob.glob(raydp_agent_path)
+        if raydp_agent_jars:
+            self._configs[SPARK_JAVAAGENT] = raydp_agent_jars[0]
         # for JVM running in ray
         self._configs[SPARK_RAY_LOG4J_FACTORY_CLASS_KEY] = versions.RAY_LOG4J_VERSION
 
diff --git a/python/raydp/tf/estimator.py b/python/raydp/tf/estimator.py
index 2c4977da..808aae68 100644
--- a/python/raydp/tf/estimator.py
+++ b/python/raydp/tf/estimator.py
@@ -43,7 +43,7 @@ def __init__(self,
                  metrics: Union[List[keras.metrics.Metric], List[str]] = None,
                  feature_columns: Union[str, List[str]] = None,
                  label_columns: Union[str, List[str]] = None,
-                 merge_feature_columns: bool = False,
+                 merge_feature_columns: bool = True,
                  batch_size: int = 128,
                  drop_last: bool = False,
                  num_epochs: int = 1,
@@ -268,4 +268,4 @@ def fit_on_spark(self,
 
     def get_model(self) -> Any:
         assert self._trainer, "Trainer has not been created"
-        return TensorflowCheckpoint(self._results.checkpoint).get_model()
+        return TensorflowCheckpoint(self._results.checkpoint.to_directory()).get_model()
diff --git a/python/raydp/torch/estimator.py b/python/raydp/torch/estimator.py
index 6b3440a7..c03d3602 100644
--- a/python/raydp/torch/estimator.py
+++ b/python/raydp/torch/estimator.py
@@ -378,4 +378,4 @@ def fit_on_spark(self,
 
     def get_model(self):
         assert self._trainer is not None, "Must call fit first"
-        return TorchCheckpoint(self._trained_results.checkpoint.as_directory()).get_model()
+        return TorchCheckpoint(self._trained_results.checkpoint.to_directory()).get_model()
diff --git a/python/raydp/xgboost/estimator.py b/python/raydp/xgboost/estimator.py
index 120b76a6..71e48ffb 100644
--- a/python/raydp/xgboost/estimator.py
+++ b/python/raydp/xgboost/estimator.py
@@ -109,4 +109,4 @@ def fit_on_spark(self,
             train_ds, evaluate_ds, max_retries)
 
     def get_model(self):
-        return XGBoostCheckpoint.from_checkpoint(self._results.checkpoint).get_model()
+        return XGBoostCheckpoint(self._results.checkpoint.to_directory()).get_model()