Fix callback issue in can (#265)

Author: YanisLalou

skada/deep/_divergence.py

@@ -255,6 +255,7 @@ def CAN(
     class_threshold=3,
     sigmas=None,
     base_criterion=None,
+    callbacks=None,
     **kwargs,
 ):
     """Contrastive Adaptation Network (CAN) domain adaptation method.
@@ -281,6 +282,8 @@ def CAN(
     base_criterion : torch criterion (class)
         The base criterion used to compute the loss with source
         labels. If None, the default is `torch.nn.CrossEntropyLoss`.
+    callbacks : list, optional
+        List of callbacks to be used during training.
 
     References
     ----------
@@ -292,6 +295,14 @@ def CAN(
     if base_criterion is None:
         base_criterion = torch.nn.CrossEntropyLoss()
 
+    if callbacks is None:
+        callbacks = [ComputeSourceCentroids()]
+    else:
+        if isinstance(callbacks, list):
+            callbacks.append(ComputeSourceCentroids())
+        else:
+            callbacks = [callbacks, ComputeSourceCentroids()]
+
     net = DomainAwareNet(
         module=DomainAwareModule,
         module__base_module=module,
@@ -305,7 +316,7 @@ def CAN(
             class_threshold=class_threshold,
             sigmas=sigmas,
         ),
-        callbacks=[ComputeSourceCentroids()],
+        callbacks=callbacks,
         **kwargs,
     )
     return net

skada/deep/tests/test_deep_divergence.py

@@ -4,6 +4,7 @@
 #
 # License: BSD 3-Clause
 import pytest
+from skorch.callbacks import EpochScoring
 
 pytest.importorskip("torch")
 
@@ -149,3 +150,48 @@ def test_can(sigmas, distance_threshold, class_threshold):
 
     history = method.history_
     assert history[0]["train_loss"] > history[-1]["train_loss"]
+
+
+def test_can_with_custom_callbacks():
+    module = ToyModule2D()
+    module.eval()
+
+    n_samples = 10
+    dataset = make_shifted_datasets(
+        n_samples_source=n_samples,
+        n_samples_target=n_samples,
+        shift="concept_drift",
+        noise=0.1,
+        random_state=42,
+        return_dataset=True,
+    )
+
+    # Create a custom callback
+    custom_callback = EpochScoring(scoring="accuracy", lower_is_better=False)
+
+    method = CAN(
+        ToyModule2D(),
+        reg=0.01,
+        layer_name="dropout",
+        batch_size=10,
+        max_epochs=10,
+        train_split=None,
+        callbacks=[custom_callback],  # Pass the custom callback
+    )
+
+    X, y, sample_domain = dataset.pack_train(as_sources=["s"], as_targets=["t"])
+    method.fit(X.astype(np.float32), y, sample_domain)
+
+    X_test, y_test, sample_domain_test = dataset.pack_test(as_targets=["t"])
+
+    y_pred = method.predict(X_test.astype(np.float32), sample_domain_test)
+
+    assert y_pred.shape[0] == X_test.shape[0]
+
+    history = method.history_
+    assert history[0]["train_loss"] > history[-1]["train_loss"]
+
+    # Check if both custom callback and ComputeSourceCentroids are present
+    callback_classes = [cb.__class__.__name__ for cb in method.callbacks]
+    assert "EpochScoring" in callback_classes
+    assert "ComputeSourceCentroids" in callback_classes