🐛 Fix lr scale #59

Co-authored-by: hanruisong00 <[email protected]> Co-authored-by: Adrien Lafage <[email protected]>
ENSTA-U2IS-AI · Oct 27, 2023 · d6df25c · d6df25c
1 parent 5e67064
commit d6df25c
Showing 1 changed file with 3 additions and 0 deletions.
diff --git 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