Merge branch 'dev' of github.com:ENSTA-U2IS/torch-uncertainty into vi…

…sualization

torch_uncertainty/routines/classification.py

@@ -456,11 +456,27 @@ def training_step(
         self, batch: Tuple[Tensor, Tensor], batch_idx: int
     ) -> STEP_OUTPUT:
         batch = self.mixup(*batch)
+        # eventual input repeat is done in the model
 
+        if self.is_elbo:
+            loss = self.criterion(inputs, targets)
+        else:
+            logits = self.forward(inputs)
+            # BCEWithLogitsLoss expects float targets
+            if self.binary_cls and self.loss == nn.BCEWithLogitsLoss:
+                logits = logits.squeeze(-1)
+                targets = targets.float()
         # eventual input repeat is done in the model
         inputs, targets = self.format_batch_fn(batch)
         logits = self.forward(inputs)
 
+        # BCEWithLogitsLoss expects float targets
+        if self.binary_cls and self.loss == nn.BCEWithLogitsLoss:
+            logits = logits.squeeze(-1)
+            targets = targets.float()
+
+            loss = self.criterion(logits, targets)
+
         self.log("train_loss", loss)
         return loss