From 2ad877dad75ac60e77d665c6c39d9221b989f7f0 Mon Sep 17 00:00:00 2001
From: Anton Baumann <anton@antonbaumann.com>
Date: Thu, 2 May 2024 14:35:24 +0200
Subject: [PATCH] improve documentation

---
 utools/wrappers/__init__.py    | 4 ++--
 utools/wrappers/ensemble.py    | 4 ++--
 utools/wrappers/monte_carlo.py | 6 +++---
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/utools/wrappers/__init__.py b/utools/wrappers/__init__.py
index 1b74d1e..0d10b2c 100644
--- a/utools/wrappers/__init__.py
+++ b/utools/wrappers/__init__.py
@@ -3,7 +3,7 @@
 from utools.wrappers.ensemble import Ensemble
 
 __all__ = [
-    'BaseWrapper',
-    'MonteCarlo',
     'Ensemble',
+    'MonteCarlo',
+    'BaseWrapper',
 ]
\ No newline at end of file
diff --git a/utools/wrappers/ensemble.py b/utools/wrappers/ensemble.py
index 649e6c2..1daf6f8 100644
--- a/utools/wrappers/ensemble.py
+++ b/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)
diff --git a/utools/wrappers/monte_carlo.py b/utools/wrappers/monte_carlo.py
index fa549c7..11044df 100644
--- a/utools/wrappers/monte_carlo.py
+++ b/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):