fix: problem of type with BCEwithLogitsLoss

Author: meilame-tayebjee

torchTextClassifiers/dataset/dataset.py

@@ -100,7 +100,7 @@ def collate_fn(self, batch):
                 labels_tensor[rows, cols] = 1
 
             else:
-                labels_tensor = torch.tensor(labels, dtype=torch.long)
+                labels_tensor = torch.tensor(labels)
         else:
             labels_tensor = None
 

torchTextClassifiers/model/lightning.py

@@ -76,6 +76,10 @@ def training_step(self, batch, batch_idx: int) -> torch.Tensor:
         targets = batch["labels"]
 
         outputs = self.forward(batch)
+
+        if isinstance(self.loss, torch.nn.BCEWithLogitsLoss):
+            targets = targets.float()
+
         loss = self.loss(outputs, targets)
         self.log("train_loss", loss, on_epoch=True, on_step=True, prog_bar=True)
         accuracy = self.accuracy_fn(outputs, targets)