diff --git a/Dockerfile b/Dockerfile
index 4735a11e9b..e2a054531d 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -85,7 +85,7 @@ ARG WITH_SQLFLOW_MODELS="ON"
 RUN if [ "${WITH_SQLFLOW_MODELS:-ON}" = "ON" ]; then \
   git clone https://github.com/sql-machine-learning/models.git && \
   cd models && \
-  git checkout 58f4c137129e2bc749320bafcc8fddb7c737fed9 && \
+  git checkout a3559618a013820385f43307261ad34351da2fbf && \
   bash -c "source activate sqlflow-dev && python setup.py install" && \
   cd .. && \
   rm -rf models; \
diff --git a/cmd/sqlflowserver/main_test.go b/cmd/sqlflowserver/main_test.go
index 26726d26e0..9168355c78 100644
--- a/cmd/sqlflowserver/main_test.go
+++ b/cmd/sqlflowserver/main_test.go
@@ -290,11 +290,11 @@ func TestEnd2EndMySQL(t *testing.T) {
 	// Cases using feature derivation
 	t.Run("CaseTrainTextClassificationIR", CaseTrainTextClassificationIR)
 	t.Run("CaseTrainTextClassificationFeatureDerivation", CaseTrainTextClassificationFeatureDerivation)
-	t.Run("CaseXgboostFeatureDerivation", CaseXgboostFeatureDerivation)
+	t.Run("CaseXGBoostFeatureDerivation", CaseXGBoostFeatureDerivation)
 	t.Run("CaseTrainFeatureDerevation", CaseTrainFeatureDerevation)
 }
 
-func CaseXgboostFeatureDerivation(t *testing.T) {
+func CaseXGBoostFeatureDerivation(t *testing.T) {
 	a := assert.New(t)
 	trainSQL := `SELECT * FROM housing.train
 TO TRAIN xgboost.gbtree
diff --git a/python/sqlflow_submitter/tensorflow/train.py b/python/sqlflow_submitter/tensorflow/train.py
index 16e09ca992..9ea1a1e9aa 100644
--- a/python/sqlflow_submitter/tensorflow/train.py
+++ b/python/sqlflow_submitter/tensorflow/train.py
@@ -96,14 +96,14 @@ def validate_input_fn(batch_size):
         return dataset.batch(batch_size)
 
     if is_keras_model:
-        classifier.compile(optimizer=classifier.default_optimizer(),
-            loss=classifier.default_loss(),
+        classifier.compile(optimizer=classifier.optimizer(),
+            loss=classifier.loss(),
             metrics=["accuracy"])
         if hasattr(classifier, 'sqlflow_train_loop'):
             classifier.sqlflow_train_loop(train_input_fn(batch_size))
         else:
             classifier.fit(train_input_fn(batch_size),
-                epochs=epochs if epochs else classifier.default_training_epochs(),
+                epochs=epochs if epochs else 1,
                 verbose=verbose)
         classifier.save_weights(save, save_format="h5")
         if label_meta["feature_name"] != "" and validate_select != "":