From d6df25c2ece40f222326695bdf846ab5178f8b44 Mon Sep 17 00:00:00 2001
From: Olivier <olivier.ar.laurent@gmail.com>
Date: Fri, 27 Oct 2023 09:51:18 +0200
Subject: [PATCH] :bug: Fix lr scale #59

Co-authored-by: hanruisong00 <hanruisong@tju.edu.cn>
Co-authored-by: Adrien Lafage <adrienlafage@outlook.com>
---
 torch_uncertainty/routines/classification.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/torch_uncertainty/routines/classification.py b/torch_uncertainty/routines/classification.py
index 0277d1f5..4df32579 100644
--- a/torch_uncertainty/routines/classification.py
+++ b/torch_uncertainty/routines/classification.py
@@ -481,6 +481,9 @@ def training_step(
             else:
                 loss = self.criterion(logits, targets, self.current_epoch)
 
+            # compensate for mean on the estimators
+            loss *= self.num_estimators
+
         self.log("train_loss", loss)
         return loss