these changes could be useful

src/xgboost_distribution/model.py

@@ -219,9 +219,12 @@ def predict(
         )
         return self._distribution.predict(params)
 
-    def virtual_ensemble_predict(self):
+    def virtual_ensemble_predict(self, X, virtual_model_count: int = 10):
         """Predict using virtual ensembles."""
 
+        params = self.predict(X, iteration_range=(0, self.best_iteration + 1))
+        return params
+
     def save_model(self, fname: Union[str, os.PathLike]) -> None:
         # self._distribution class cannot be saved by `super().save_model`, as it
         # attempts to call `json.dumps({"_distribution": self._distribution})`