Fix handling of batches without positives

Author: MGlauer

chebai/models/base.py

@@ -303,7 +303,7 @@ def _execute(
                     logger=True,
                     sync_dist=sync_dist,
                 )
-            if metrics and labels is not None:
+            if metrics and labels is not None and pr is not None:
                 for metric_name, metric in metrics.items():
                     metric.update(pr, tar)
                 self._log_metrics(prefix, metrics, len(batch))

chebai/models/electra.py

@@ -333,7 +333,10 @@ def _get_prediction_and_labels(
                 missing_labels = loss_kwargs["missing_labels"]
                 d[missing_labels] = 0
             has_positive_entries = torch.sum(labels, dim=-1)>0
-            return d[has_positive_entries], labels[has_positive_entries].int() if labels[has_positive_entries] is not None else None
+            if torch.sum(has_positive_entries):
+                return d[has_positive_entries], labels[has_positive_entries].int() if labels[has_positive_entries] is not None else None
+            else:
+                return None, None
         elif self.model_type == 'regression':
             return d, labels
         else: