update function signature

Author: xuyxu

torchensemble/_constants.py

@@ -174,7 +174,8 @@
 
 
 __vectorize_doc = """
-    Return the vectorization result of the ensemble using functorch.
+    Return the vectorization result of the ensemble using functorch. Details
+    available at `functorch model ensembling <https://pytorch.org/functorch/stable/notebooks/ensembling.html>`_.
 
     Returns
     -------

torchensemble/voting.py

@@ -307,6 +307,10 @@ def evaluate(self, test_loader, return_loss=False):
     def predict(self, *x):
         return super().predict(*x)
 
+    @torchensemble_model_doc(item="vectorize")
+    def vectorize(self):
+        return super().vectorize()
+
 
 @torchensemble_model_doc(
     """Implementation on the NeuralForestClassifier.""", "tree_ensemble_model"
@@ -374,6 +378,10 @@ def fit(
             save_dir=save_dir,
         )
 
+    @torchensemble_model_doc(item="vectorize")
+    def vectorize(self):
+        return super().vectorize()
+
 
 @torchensemble_model_doc("""Implementation on the VotingRegressor.""", "model")
 class VotingRegressor(BaseRegressor):
@@ -559,6 +567,10 @@ def evaluate(self, test_loader):
     def predict(self, *x):
         return super().predict(*x)
 
+    @torchensemble_model_doc(item="vectorize")
+    def vectorize(self):
+        return super().vectorize()
+
 
 @torchensemble_model_doc(
     """Implementation on the NeuralForestRegressor.""", "tree_ensemble_model"
@@ -620,3 +632,7 @@ def fit(
             save_model=save_model,
             save_dir=save_dir,
         )
+
+    @torchensemble_model_doc(item="vectorize")
+    def vectorize(self):
+        return super().vectorize()