fix: conditional loss breaking for batch size one (#60)

rul_adapt/approach/pseudo_labels.py

@@ -65,6 +65,8 @@ def generate_pseudo_labels(
     values. It is recommended to clip them to zero and `max_rul` respectively before
     using them to patch a reader.
 
+    The model is assumed to reside on the CPU where the calculation will be performed.
+
     Args:
         dm: The data module to generate pseudo labels for.
         model: The model to use for generating the pseudo labels.

rul_adapt/loss/conditional.py

@@ -107,7 +107,7 @@ def compute(self) -> torch.Tensor:
 
 
 def _membership(preds: torch.Tensor, fuzzy_set: Tuple[float, float]) -> torch.Tensor:
-    preds = preds.squeeze() if len(preds.shape) > 1 else preds
+    preds = preds.squeeze(-1) if preds.ndim > 1 else preds
     membership = (preds >= fuzzy_set[0]) & (preds < fuzzy_set[1])
 
     return membership

tests/test_loss/test_conditional.py

@@ -115,6 +115,18 @@ def test__membership():
     assert torch.all(_membership(inputs, fuzzy_set) == expected)
 
 
+@pytest.mark.parametrize("loss_fixture", ["cdann", "cmmd"])
+def test_forward_batch_size_one(loss_fixture, request):
+    """Should not fail for batch size of one."""
+    loss_func = request.getfixturevalue(loss_fixture)
+    source = torch.rand(1, 10)
+    source_preds = torch.zeros(1, 1)
+    target = torch.rand(1, 10)
+    target_preds = torch.zeros(1, 1)
+
+    loss_func(source, source_preds, target, target_preds)
+
+
 def test_backward_cdann(cdann):
     source = torch.rand(10, 10)
     source_preds = torch.zeros(10, 1)