diff --git a/tests/unit/aiplatform/test_metadata_models.py b/tests/unit/aiplatform/test_metadata_models.py
index b7120c2310..6d0fedcda8 100644
--- a/tests/unit/aiplatform/test_metadata_models.py
+++ b/tests/unit/aiplatform/test_metadata_models.py
@@ -18,6 +18,7 @@
 import datetime
 import os
 import pickle
+import sys
 from importlib import reload
 from unittest import mock
 from unittest.mock import patch
@@ -41,6 +42,7 @@
 from sklearn.linear_model import LinearRegression
 import tensorflow as tf
 import xgboost as xgb
+import xgboost.sklearn
 
 
 # project
@@ -339,6 +341,8 @@ def setup_method(self):
         reload(initializer)
         reload(metadata)
         reload(aiplatform)
+        reload(xgboost.sklearn)
+        reload(xgb)
 
     def teardown_method(self):
         initializer.global_pool.shutdown(wait=True)