🎨 Update typing in _RegDeepEnsembles

ENSTA-U2IS-AI · Nov 28, 2024 · 9d4ec19 · 9d4ec19
1 parent 11afd5b
commit 9d4ec19
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_uncertainty/models/wrappers/deep_ensembles.py b/torch_uncertainty/models/wrappers/deep_ensembles.py
@@ -40,7 +40,7 @@ def __init__(
         super().__init__(models)
         self.probabilistic = probabilistic
 
-    def forward(self, x: torch.Tensor) -> Distribution:
+    def forward(self, x: torch.Tensor) -> torch.Tensor | dict[str, torch.Tensor]:
         r"""Return the logits of the ensemble.
 
         Args: