Skip to content

Commit 664693a

Browse files
sararobcopybara-github
authored andcommitted
chore: Make XGB tests compatible with sklearn 1.8
PiperOrigin-RevId: 858666892
1 parent 682e65f commit 664693a

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

tests/unit/aiplatform/test_metadata_models.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,7 @@ def mock_storage_blob_download_xgboost_xgbmodel_file():
159159
def create_model_file(filename):
160160
x, y = make_classification()
161161
model = xgb.XGBClassifier()
162+
model._estimator_type = "classifier"
162163
model.fit(x, y)
163164
model.save_model(filename)
164165

@@ -522,6 +523,7 @@ def test_save_model_xgboost_xgbmodel(
522523

523524
x, y = make_classification()
524525
xgb_model = xgb.XGBClassifier()
526+
xgb_model._estimator_type = "classifier"
525527
xgb_model.fit(x, y)
526528

527529
aiplatform.init(

0 commit comments

Comments
 (0)