improve documentation

utools/wrappers/__init__.py

@@ -3,7 +3,7 @@
 from utools.wrappers.ensemble import Ensemble
 
 __all__ = [
-    'BaseWrapper',
-    'MonteCarlo',
     'Ensemble',
+    'MonteCarlo',
+    'BaseWrapper',
 ]

utools/wrappers/ensemble.py

@@ -25,7 +25,7 @@ def __init__(
             criterion (RegressionLoss | HeteroscedasticSoftmax): The criterion to be used for computing probabilistic outputs.
         """
         super(Ensemble, self).__init__()
-        self.wrapper = BaseWrapper(models=models, criterion=criterion, monte_carlo_samples=1)
+        self.__wrapper = BaseWrapper(models=models, criterion=criterion, monte_carlo_samples=1)
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         """
@@ -37,4 +37,4 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         Returns:
             torch.Tensor: The ensemble prediction dictionary. 
         """
-        return self.wrapper(input)
+        return self.__wrapper(input)

utools/wrappers/monte_carlo.py

@@ -26,17 +26,17 @@ def __init__(
             monte_carlo_samples (int): The number of samples to use for Monte Carlo simulation.
         """
         super(MonteCarlo, self).__init__()
-        self.wrapper = BaseWrapper(models=[model], criterion=criterion, monte_carlo_samples=monte_carlo_samples)
+        self.__wrapper = BaseWrapper(models=[model], criterion=criterion, monte_carlo_samples=monte_carlo_samples)
 
         # Activate MC Dropout for the model
-        for model in self.wrapper.models:
+        for model in self.__wrapper.models:
             self._activate_mc_dropout(model)
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         """
         Computes the Monte Carlo prediction.
         """
-        return self.wrapper(input)
+        return self.__wrapper(input)
 
     @staticmethod
     def _activate_mc_dropout(model: torch.nn.Module):