diff --git a/python-package/lightgbm/basic.py b/python-package/lightgbm/basic.py
index 41cb24650788..c736fedfb9ee 100644
--- a/python-package/lightgbm/basic.py
+++ b/python-package/lightgbm/basic.py
@@ -1736,7 +1736,7 @@ def _set_init_score_by_predictor(
         self,
         predictor: Optional[_InnerPredictor],
         data: _LGBM_TrainDataType,
-        used_indices: Optional[List[int]]
+        used_indices: Optional[Union[List[int], np.ndarray]]
     ) -> "Dataset":
         data_has_header = False
         if isinstance(data, (str, Path)) and self.params is not None: