From 6b8ee399f3dc6ac6915bd58a97177033992cddec Mon Sep 17 00:00:00 2001
From: mbarneche <aruket.ipa@gmail.com>
Date: Tue, 10 Jun 2025 17:33:25 +0200
Subject: [PATCH 01/19] Create _test_time.py

---
 skada/deep/_test_time.py | 128 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 128 insertions(+)
 create mode 100644 skada/deep/_test_time.py

diff --git a/skada/deep/_test_time.py b/skada/deep/_test_time.py
new file mode 100644
index 00000000..d08943b2
--- /dev/null
+++ b/skada/deep/_test_time.py
@@ -0,0 +1,128 @@
+# Author : Maxence Barneche
+#
+# License: BSD-3-Clause
+
+import torch
+
+from skada.deep.base import DomainAwareNet
+
+
+class TestTimeCriterion(torch.nn.Module):
+    def __init__(
+        self,
+        base_criterion,
+        adapt_criterion,
+        reg=1,
+        reduction="mean",
+        train_on_target=False,
+        mode="finetune",
+    ):
+        super().__init__()
+        self.base_criterion = base_criterion
+        self.adapt_criterion = adapt_criterion
+        self.reg = reg
+        self.train_on_target = train_on_target
+        self.mode = mode
+
+        # Update the reduce parameter for both criteria if specified
+        if hasattr(self.base_criterion, "reduction"):
+            self.base_criterion.reduction = reduction
+
+        # TODO: implement losses between source and target
+        # that are sum of the losses of each sample
+        # if hasattr(self.adapt_criterion, 'reduction'):
+        #     self.adapt_criterion.reduction = reduction
+
+    def forward(
+        self,
+        y_pred,
+        y_true,
+    ):
+        """
+        Parameters
+        ----------
+        y_pred : tuple
+            This tuple comprises all the different data
+            needed to compute DA loss:
+                - y_pred : prediction of the source and target domains
+                - domain_pred : prediction of domain classifier if given
+                - features :  features of the chosen layer
+                  of source and target domains
+                - sample_domain : giving the domain of each samples
+        y_true :
+            The true labels. Available for source, masked for target.
+        """
+        y_pred, domain_pred, features, sample_domain, sample_idx = y_pred
+        source_idx = sample_domain >= 0
+        y_pred_s = y_pred[source_idx]
+        y_pred_t = y_pred[~source_idx]
+
+        if domain_pred is not None:
+            domain_pred_s = domain_pred[source_idx]
+            domain_pred_t = domain_pred[~source_idx]
+        else:
+            domain_pred_s = None
+            domain_pred_t = None
+
+        if features is not None:
+            features_s = features[source_idx]
+            features_t = features[~source_idx]
+        else:
+            features_s = None
+            features_t = None
+
+        if sample_idx is not None:
+            sample_idx_s = sample_idx[source_idx]
+            sample_idx_t = sample_idx[~source_idx]
+        else:
+            sample_idx_s = None
+            sample_idx_t = None
+
+        if self.train_on_target:
+            base_loss = self.base_criterion(y_pred_t, y_true[~source_idx])
+        else:
+            base_loss = self.base_criterion(y_pred_s, y_true[source_idx])
+
+        if self.mode == "finetune":
+            # In finetune mode, we only compute the base loss
+            return base_loss
+        elif self.mode == "adapt":
+            self.adapt_criterion(
+                y_s=y_true[source_idx],
+                y_pred_s=y_pred_s,
+                y_pred_t=y_pred_t,
+                domain_pred_s=domain_pred_s,
+                domain_pred_t=domain_pred_t,
+                features_s=features_s,
+                features_t=features_t,
+                sample_idx_s=sample_idx_s,
+                sample_idx_t=sample_idx_t,
+            )
+
+
+class TestTimeNet(DomainAwareNet):
+    def __init__(self, module, criterion: "TestTimeCriterion", **kwargs):
+        super().__init__(module, criterion=criterion, **kwargs)
+
+    def fit_prepare(
+        self, X, y=None, sample_domain=None, sample_weight=None, **fit_params
+    ):
+        X = self._prepare_input(X, y, sample_domain, sample_weight)
+        X = X.select_source()
+        self.partial_fit(X, None, **fit_params)
+        return self
+
+    def fit_adapt(self, X, sample_domain=None, sample_weight=None, **fit_params):
+        if sample_domain is None:
+            sample_domain = -1
+        X = self._prepare_input(X, None, sample_domain, sample_weight)
+        X = X.select_target()
+        self.criterion.mode = "adapt"
+        self.partial_fit(X, None, **fit_params)
+        # here should be a partial fit call with criterion_adapt
+
+    def fit(self, X, y=None, sample_domain=None, sample_weight=None, **fit_params):
+        X = self._prepare_input(X, y, sample_domain, sample_weight)
+        self.fit_prepare(X, **fit_params)
+        self.fit_adapt(X, **fit_params)
+        return self

From cd13eed1233c0e45d1cca957c73a62831d5a63e0 Mon Sep 17 00:00:00 2001
From: mbarneche <aruket.ipa@gmail.com>
Date: Wed, 11 Jun 2025 12:06:23 +0200
Subject: [PATCH 02/19] removed unnecessary clauses in TestTimeCriterion

---
 skada/deep/_test_time.py | 54 ++++++++--------------------------------
 1 file changed, 10 insertions(+), 44 deletions(-)

diff --git a/skada/deep/_test_time.py b/skada/deep/_test_time.py
index d08943b2..c4b616c1 100644
--- a/skada/deep/_test_time.py
+++ b/skada/deep/_test_time.py
@@ -15,14 +15,12 @@ def __init__(
         reg=1,
         reduction="mean",
         train_on_target=False,
-        mode="finetune",
     ):
         super().__init__()
         self.base_criterion = base_criterion
         self.adapt_criterion = adapt_criterion
         self.reg = reg
         self.train_on_target = train_on_target
-        self.mode = mode
 
         # Update the reduce parameter for both criteria if specified
         if hasattr(self.base_criterion, "reduction"):
@@ -52,51 +50,19 @@ def forward(
         y_true :
             The true labels. Available for source, masked for target.
         """
-        y_pred, domain_pred, features, sample_domain, sample_idx = y_pred
-        source_idx = sample_domain >= 0
-        y_pred_s = y_pred[source_idx]
-        y_pred_t = y_pred[~source_idx]
-
-        if domain_pred is not None:
-            domain_pred_s = domain_pred[source_idx]
-            domain_pred_t = domain_pred[~source_idx]
-        else:
-            domain_pred_s = None
-            domain_pred_t = None
-
-        if features is not None:
-            features_s = features[source_idx]
-            features_t = features[~source_idx]
-        else:
-            features_s = None
-            features_t = None
-
-        if sample_idx is not None:
-            sample_idx_s = sample_idx[source_idx]
-            sample_idx_t = sample_idx[~source_idx]
-        else:
-            sample_idx_s = None
-            sample_idx_t = None
+        y_pred, domain_pred, features, _, sample_idx = y_pred
 
         if self.train_on_target:
-            base_loss = self.base_criterion(y_pred_t, y_true[~source_idx])
-        else:
-            base_loss = self.base_criterion(y_pred_s, y_true[source_idx])
-
-        if self.mode == "finetune":
             # In finetune mode, we only compute the base loss
-            return base_loss
-        elif self.mode == "adapt":
-            self.adapt_criterion(
-                y_s=y_true[source_idx],
-                y_pred_s=y_pred_s,
-                y_pred_t=y_pred_t,
-                domain_pred_s=domain_pred_s,
-                domain_pred_t=domain_pred_t,
-                features_s=features_s,
-                features_t=features_t,
-                sample_idx_s=sample_idx_s,
-                sample_idx_t=sample_idx_t,
+            return self.base_criterion(y_pred, y_true)
+        else:
+            # In adapt mode, we compute the adaptation loss
+            return self.adapt_criterion(
+                y_s=y_true,
+                y_pred=y_pred,
+                domain_pred=domain_pred,
+                features=features,
+                sample_idx=sample_idx,
             )
 
 

From d0ee0a3d7d4baed1bffbc811b35ecf1b17f92864 Mon Sep 17 00:00:00 2001
From: mbarneche <aruket.ipa@gmail.com>
Date: Thu, 12 Jun 2025 10:30:16 +0200
Subject: [PATCH 03/19] added tests for TestTimeNet

---
 skada/deep/tests/test_test_time.py | 39 ++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)
 create mode 100644 skada/deep/tests/test_test_time.py

diff --git a/skada/deep/tests/test_test_time.py b/skada/deep/tests/test_test_time.py
new file mode 100644
index 00000000..df175f28
--- /dev/null
+++ b/skada/deep/tests/test_test_time.py
@@ -0,0 +1,39 @@
+# Author : Maxence Barneche
+#
+# License: BSD-3-Clause
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from skada.datasets import make_shifted_datasets
+from skada.deep._test_time import TestTimeCriterion, TestTimeNet
+from skada.deep.base import DomainAwareModule, DomainBalancedDataLoader
+from skada.deep.losses import TestLoss
+from skada.deep.modules import ToyModule2D
+
+
+def test_test_time_criterion():
+    num_features = 10
+    module = ToyModule2D(num_features=num_features)
+    criterion = TestTimeCriterion(torch.nn.CrossEntropyLoss(), TestLoss())
+
+    n_samples = 20
+    dataset = make_shifted_datasets(
+        n_samples_source=n_samples,
+        n_samples_target=n_samples,
+        shift="conditional_shift",
+        noise=0.1,
+        random_state=42,
+    )
+    X, y, sample_domain = dataset
+    method = TestTimeNet(
+        DomainAwareModule(module, "dropout"),
+        iterator_train=DomainBalancedDataLoader,
+        criterion=criterion,
+        batch_size=10,
+        max_epochs=2,
+        train_split=None,
+    )
+
+    method.fit(X, y, sample_domain)

From 32aec72ac844afab6a46719289a746d7a17dfbfc Mon Sep 17 00:00:00 2001
From: mbarneche <aruket.ipa@gmail.com>
Date: Thu, 12 Jun 2025 10:35:36 +0200
Subject: [PATCH 04/19] Corrected wrong attribute call

---
 skada/deep/_test_time.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/skada/deep/_test_time.py b/skada/deep/_test_time.py
index c4b616c1..d649caa9 100644
--- a/skada/deep/_test_time.py
+++ b/skada/deep/_test_time.py
@@ -75,7 +75,10 @@ def fit_prepare(
     ):
         X = self._prepare_input(X, y, sample_domain, sample_weight)
         X = X.select_source()
+        temp = self.criterion.train_on_target
+        self.criterion.train_on_target = False
         self.partial_fit(X, None, **fit_params)
+        self.criterion.train_on_target = temp
         return self
 
     def fit_adapt(self, X, sample_domain=None, sample_weight=None, **fit_params):
@@ -83,9 +86,11 @@ def fit_adapt(self, X, sample_domain=None, sample_weight=None, **fit_params):
             sample_domain = -1
         X = self._prepare_input(X, None, sample_domain, sample_weight)
         X = X.select_target()
-        self.criterion.mode = "adapt"
+        temp = self.criterion.train_on_target
+        self.criterion.train_on_target = True
         self.partial_fit(X, None, **fit_params)
-        # here should be a partial fit call with criterion_adapt
+        self.criterion.train_on_target = temp
+        return self
 
     def fit(self, X, y=None, sample_domain=None, sample_weight=None, **fit_params):
         X = self._prepare_input(X, y, sample_domain, sample_weight)

From c23e8f271b9942687a8d038d7c0f3a741a9ccb56 Mon Sep 17 00:00:00 2001
From: mbarneche <aruket.ipa@gmail.com>
Date: Thu, 12 Jun 2025 11:05:27 +0200
Subject: [PATCH 05/19] fixed wrong dataloader in test

---
 skada/deep/tests/test_test_time.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/skada/deep/tests/test_test_time.py b/skada/deep/tests/test_test_time.py
index df175f28..a883deb1 100644
--- a/skada/deep/tests/test_test_time.py
+++ b/skada/deep/tests/test_test_time.py
@@ -6,9 +6,11 @@
 
 torch = pytest.importorskip("torch")
 
+from torch.utils.data import DataLoader
+
 from skada.datasets import make_shifted_datasets
 from skada.deep._test_time import TestTimeCriterion, TestTimeNet
-from skada.deep.base import DomainAwareModule, DomainBalancedDataLoader
+from skada.deep.base import DomainAwareModule
 from skada.deep.losses import TestLoss
 from skada.deep.modules import ToyModule2D
 
@@ -29,7 +31,7 @@ def test_test_time_criterion():
     X, y, sample_domain = dataset
     method = TestTimeNet(
         DomainAwareModule(module, "dropout"),
-        iterator_train=DomainBalancedDataLoader,
+        iterator_train=DataLoader,
         criterion=criterion,
         batch_size=10,
         max_epochs=2,

From a3c0ced280a12af6e5cc7863f0374fdba28758c1 Mon Sep 17 00:00:00 2001
From: mbarneche <aruket.ipa@gmail.com>
Date: Thu, 12 Jun 2025 16:07:41 +0200
Subject: [PATCH 06/19] remade tests, improved TestTimeNet

---
 skada/deep/_test_time.py                      | 57 ++++++++++---------
 skada/deep/base.py                            |  4 +-
 skada/deep/losses.py                          |  4 +-
 ...st_test_time.py => test_deep_test_time.py} | 18 +++++-
 4 files changed, 54 insertions(+), 29 deletions(-)
 rename skada/deep/tests/{test_test_time.py => test_deep_test_time.py} (62%)

diff --git a/skada/deep/_test_time.py b/skada/deep/_test_time.py
index d649caa9..39bd4cd6 100644
--- a/skada/deep/_test_time.py
+++ b/skada/deep/_test_time.py
@@ -39,61 +39,66 @@ def forward(
         """
         Parameters
         ----------
-        y_pred : tuple
-            This tuple comprises all the different data
-            needed to compute DA loss:
-                - y_pred : prediction of the source and target domains
-                - domain_pred : prediction of domain classifier if given
-                - features :  features of the chosen layer
-                  of source and target domains
-                - sample_domain : giving the domain of each samples
+        y_pred :
+            Prediction of the labels.
         y_true :
             The true labels. Available for source, masked for target.
         """
-        y_pred, domain_pred, features, _, sample_idx = y_pred
-
         if self.train_on_target:
             # In finetune mode, we only compute the base loss
-            return self.base_criterion(y_pred, y_true)
-        else:
-            # In adapt mode, we compute the adaptation loss
             return self.adapt_criterion(
                 y_s=y_true,
                 y_pred=y_pred,
-                domain_pred=domain_pred,
-                features=features,
-                sample_idx=sample_idx,
             )
+        else:
+            # In adapt mode, we compute the adaptation loss
+            return self.base_criterion(y_pred, y_true)
 
 
 class TestTimeNet(DomainAwareNet):
-    def __init__(self, module, criterion: "TestTimeCriterion", **kwargs):
+    def __init__(
+        self,
+        module,
+        criterion: "TestTimeCriterion",
+        optimizer_adapt=None,
+        epochs_adapt=None,
+        **kwargs,
+    ):
         super().__init__(module, criterion=criterion, **kwargs)
+        self.optimizer_adapt = optimizer_adapt
+        self.epochs_adapt = epochs_adapt
 
-    def fit_prepare(
+    def fit_source(
         self, X, y=None, sample_domain=None, sample_weight=None, **fit_params
     ):
+        print("Training model on source domain...")
         X = self._prepare_input(X, y, sample_domain, sample_weight)
         X = X.select_source()
-        temp = self.criterion.train_on_target
         self.criterion.train_on_target = False
         self.partial_fit(X, None, **fit_params)
-        self.criterion.train_on_target = temp
         return self
 
     def fit_adapt(self, X, sample_domain=None, sample_weight=None, **fit_params):
+        print("Adapting model to target domain...")
         if sample_domain is None:
             sample_domain = -1
+        if self.optimizer_adapt is not None:
+            self.initialize_adapt_optimizer()
         X = self._prepare_input(X, None, sample_domain, sample_weight)
         X = X.select_target()
-        temp = self.criterion.train_on_target
         self.criterion.train_on_target = True
-        self.partial_fit(X, None, **fit_params)
-        self.criterion.train_on_target = temp
+        self.partial_fit(X, None, epochs=self.epochs_adapt, **fit_params)
         return self
 
     def fit(self, X, y=None, sample_domain=None, sample_weight=None, **fit_params):
-        X = self._prepare_input(X, y, sample_domain, sample_weight)
-        self.fit_prepare(X, **fit_params)
-        self.fit_adapt(X, **fit_params)
+        self.fit_source(X, y, sample_domain, sample_weight, **fit_params)
+        self.fit_adapt(X, sample_domain, sample_weight, **fit_params)
+        return self
+
+    def initialize_adapt_optimizer(self):
+        named_parameters = self.get_all_learnable_params()
+        args, kwargs = self.get_params_for_optimizer(
+            "adapt_optimizer", named_parameters
+        )
+        self.optimizer_ = self.optimizer_adapt(*args, **kwargs)
         return self
diff --git a/skada/deep/base.py b/skada/deep/base.py
index 40364b89..de8f4368 100644
--- a/skada/deep/base.py
+++ b/skada/deep/base.py
@@ -1032,7 +1032,9 @@ def _if_given_dict(self, d: dict, y, sample_domain, sample_weight):
         self._initialize(X, y, sample_domain, sample_weight)
 
     def _initialize(self, X:torch.Tensor, y, sample_domain, sample_weight):
-        if sample_domain is None or len(sample_domain) == 0:
+        if sample_domain is None or (
+            hasattr(sample_domain, '__len__') and len(sample_domain) == 0
+            ):
             sample_domain = _DEFAULT_SAMPLE_DOMAIN_
         if isinstance(sample_domain, int):
             sample_domain = torch.full((X.shape[0],), sample_domain)
diff --git a/skada/deep/losses.py b/skada/deep/losses.py
index 98c70c4c..1cb664df 100644
--- a/skada/deep/losses.py
+++ b/skada/deep/losses.py
@@ -346,7 +346,9 @@ def forward(
         **kwargs,
     ):
         """Compute the domain adaptation loss"""
-        return 0
+        res = torch.tensor(0, dtype=torch.float32)
+        res.requires_grad = True
+        return res
 
 
 def probability_scaling(logits, temperature=1):
diff --git a/skada/deep/tests/test_test_time.py b/skada/deep/tests/test_deep_test_time.py
similarity index 62%
rename from skada/deep/tests/test_test_time.py
rename to skada/deep/tests/test_deep_test_time.py
index a883deb1..e432e44c 100644
--- a/skada/deep/tests/test_test_time.py
+++ b/skada/deep/tests/test_deep_test_time.py
@@ -6,6 +6,8 @@
 
 torch = pytest.importorskip("torch")
 
+import numpy as np
+from torch.optim import Adam
 from torch.utils.data import DataLoader
 
 from skada.datasets import make_shifted_datasets
@@ -15,7 +17,11 @@
 from skada.deep.modules import ToyModule2D
 
 
-def test_test_time_criterion():
+@pytest.mark.parametrize(
+    "sd, epochs_adapt, optimizer_adapt",
+    [(True, 3, Adam), (False, None, None)],
+)
+def test_test_time_criterion(sd, epochs_adapt, optimizer_adapt):
     num_features = 10
     module = ToyModule2D(num_features=num_features)
     criterion = TestTimeCriterion(torch.nn.CrossEntropyLoss(), TestLoss())
@@ -29,13 +35,23 @@ def test_test_time_criterion():
         random_state=42,
     )
     X, y, sample_domain = dataset
+    X = X.astype(np.float32)
+    sample_domain = sample_domain if sd else None
+
     method = TestTimeNet(
         DomainAwareModule(module, "dropout"),
+        epochs_adapt=epochs_adapt,
         iterator_train=DataLoader,
         criterion=criterion,
         batch_size=10,
         max_epochs=2,
         train_split=None,
+        optimizer_adapt=optimizer_adapt,
     )
 
     method.fit(X, y, sample_domain)
+
+    # If the adaptation optimizer is specified, the net's optimizer should have changed
+    # during the fit_adapt method called inside the fit method
+    if optimizer_adapt is not None:
+        assert isinstance(method.optimizer_, optimizer_adapt)

From 28ce3f3275bcac22325c29fadd860063fb9e8d6f Mon Sep 17 00:00:00 2001
From: mbarneche <aruket.ipa@gmail.com>
Date: Thu, 12 Jun 2025 16:54:34 +0200
Subject: [PATCH 07/19] Updated tests

---
 skada/deep/_test_time.py                |  4 +---
 skada/deep/tests/test_deep_test_time.py | 11 +++++++----
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/skada/deep/_test_time.py b/skada/deep/_test_time.py
index 39bd4cd6..5b38a1de 100644
--- a/skada/deep/_test_time.py
+++ b/skada/deep/_test_time.py
@@ -78,10 +78,8 @@ def fit_source(
         self.partial_fit(X, None, **fit_params)
         return self
 
-    def fit_adapt(self, X, sample_domain=None, sample_weight=None, **fit_params):
+    def fit_adapt(self, X, sample_domain, sample_weight=None, **fit_params):
         print("Adapting model to target domain...")
-        if sample_domain is None:
-            sample_domain = -1
         if self.optimizer_adapt is not None:
             self.initialize_adapt_optimizer()
         X = self._prepare_input(X, None, sample_domain, sample_weight)
diff --git a/skada/deep/tests/test_deep_test_time.py b/skada/deep/tests/test_deep_test_time.py
index e432e44c..1b5c2003 100644
--- a/skada/deep/tests/test_deep_test_time.py
+++ b/skada/deep/tests/test_deep_test_time.py
@@ -18,10 +18,14 @@
 
 
 @pytest.mark.parametrize(
-    "sd, epochs_adapt, optimizer_adapt",
-    [(True, 3, Adam), (False, None, None)],
+    "epochs_adapt, optimizer_adapt",
+    [
+        (None, None),  # Default test
+        (3, None),  # Test with specific adaptation epochs given
+        (None, Adam),
+    ],  # Test with Adam optimizer for adaptation
 )
-def test_test_time_criterion(sd, epochs_adapt, optimizer_adapt):
+def test_test_time_criterion(epochs_adapt, optimizer_adapt):
     num_features = 10
     module = ToyModule2D(num_features=num_features)
     criterion = TestTimeCriterion(torch.nn.CrossEntropyLoss(), TestLoss())
@@ -36,7 +40,6 @@ def test_test_time_criterion(sd, epochs_adapt, optimizer_adapt):
     )
     X, y, sample_domain = dataset
     X = X.astype(np.float32)
-    sample_domain = sample_domain if sd else None
 
     method = TestTimeNet(
         DomainAwareModule(module, "dropout"),

From e81a62eafae4f33d3791669744f56e7347533660 Mon Sep 17 00:00:00 2001
From: mbarneche <aruket.ipa@gmail.com>
Date: Fri, 13 Jun 2025 11:18:18 +0200
Subject: [PATCH 08/19] added softmax_entropy, added foundation of Tent

---
 skada/deep/_test_time.py | 55 +++++++++++++++++++++++++++++++++++++++-
 skada/deep/losses.py     |  6 +++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/skada/deep/_test_time.py b/skada/deep/_test_time.py
index 5b38a1de..23538784 100644
--- a/skada/deep/_test_time.py
+++ b/skada/deep/_test_time.py
@@ -4,7 +4,7 @@
 
 import torch
 
-from skada.deep.base import DomainAwareNet
+from skada.deep.base import BaseDALoss, DomainAwareModule, DomainAwareNet
 
 
 class TestTimeCriterion(torch.nn.Module):
@@ -100,3 +100,56 @@ def initialize_adapt_optimizer(self):
         )
         self.optimizer_ = self.optimizer_adapt(*args, **kwargs)
         return self
+
+
+class TentNet(TestTimeNet):
+    def __init__(self, module, criterion, **kwargs):
+        super().__init__(module, criterion=criterion, **kwargs)
+
+    def forward(self, X, sample_domain=None, sample_weight=None):
+        """Forward pass of the model."""
+        X = self._prepare_input(X, None, sample_domain, sample_weight)
+        # Note: we are supposed to freeze some attributes of the model
+        # during the adaptation phase, but we do not do it here.
+        return self.module(X)
+
+
+class TentLoss(BaseDALoss):
+    def __init__(
+        self, reg_dist=1, reg_cl=1, base_criterion=None, target_criterion=None
+    ):
+        super().__init__()
+        self.reg_dist = reg_dist
+        self.reg_cl = reg_cl
+        self.base_criterion = base_criterion
+        self.target_criterion = target_criterion
+
+    def forward(self, y_s, y_t):
+        loss = self.target_criterion(y_t, y_s)
+        return loss + self.reg_dist * torch.mean(y_t) + self.reg_cl * torch.mean(y_s)
+
+
+def Tent(
+    module,
+    layer_name,
+    reg_dist=1,
+    reg_cl=1,
+    base_criterion=None,
+    target_criterion=None,
+    **kwargs,
+):
+    if base_criterion is None:
+        base_criterion = torch.nn.CrossEntropyLoss()
+
+    net = TentNet(
+        module=DomainAwareModule,
+        module__base_module=module,
+        module__layer_name=layer_name,
+        criterion=TestTimeCriterion,
+        criterion__base_criterion=base_criterion,
+        criterion__adapt_criterion=TentLoss(reg_dist, reg_cl, target_criterion),
+        criterion__reg=1,
+        **kwargs,
+    )
+
+    return net
diff --git a/skada/deep/losses.py b/skada/deep/losses.py
index 1cb664df..21585d7b 100644
--- a/skada/deep/losses.py
+++ b/skada/deep/losses.py
@@ -515,3 +515,9 @@ def nap_loss(features_t, y_pred_t, memory_features, memory_outputs, sample_idx_t
     classifier_loss = torch.sum(weight_ * loss_) / (torch.sum(weight_).item() + 1e-7)
 
     return classifier_loss
+
+
+def softmax_entropy(x: torch.Tensor):
+    """Compute the entropy of the softmax probabilities."""
+    loss = x.softmax(1) * x.log_softmax(1)
+    return -torch.sum(loss, axis=1)

From 4c76ba495f150196ff4c3f4db42062411458f1ec Mon Sep 17 00:00:00 2001
From: mbarneche <aruket.ipa@gmail.com>
Date: Mon, 16 Jun 2025 15:20:31 +0200
Subject: [PATCH 09/19] Updated Tent

---
 skada/deep/_test_time.py | 34 +++++++---------------------------
 1 file changed, 7 insertions(+), 27 deletions(-)

diff --git a/skada/deep/_test_time.py b/skada/deep/_test_time.py
index 23538784..2d01c9e1 100644
--- a/skada/deep/_test_time.py
+++ b/skada/deep/_test_time.py
@@ -5,6 +5,7 @@
 import torch
 
 from skada.deep.base import BaseDALoss, DomainAwareModule, DomainAwareNet
+from skada.deep.losses import softmax_entropy
 
 
 class TestTimeCriterion(torch.nn.Module):
@@ -102,52 +103,31 @@ def initialize_adapt_optimizer(self):
         return self
 
 
-class TentNet(TestTimeNet):
-    def __init__(self, module, criterion, **kwargs):
-        super().__init__(module, criterion=criterion, **kwargs)
-
-    def forward(self, X, sample_domain=None, sample_weight=None):
-        """Forward pass of the model."""
-        X = self._prepare_input(X, None, sample_domain, sample_weight)
-        # Note: we are supposed to freeze some attributes of the model
-        # during the adaptation phase, but we do not do it here.
-        return self.module(X)
-
-
 class TentLoss(BaseDALoss):
-    def __init__(
-        self, reg_dist=1, reg_cl=1, base_criterion=None, target_criterion=None
-    ):
+    def __init__(self):
         super().__init__()
-        self.reg_dist = reg_dist
-        self.reg_cl = reg_cl
-        self.base_criterion = base_criterion
-        self.target_criterion = target_criterion
 
-    def forward(self, y_s, y_t):
-        loss = self.target_criterion(y_t, y_s)
-        return loss + self.reg_dist * torch.mean(y_t) + self.reg_cl * torch.mean(y_s)
+    def forward(self, y):
+        loss = softmax_entropy(y)
+        return loss
 
 
 def Tent(
     module,
     layer_name,
-    reg_dist=1,
-    reg_cl=1,
     base_criterion=None,
-    target_criterion=None,
     **kwargs,
 ):
     if base_criterion is None:
         base_criterion = torch.nn.CrossEntropyLoss()
 
-    net = TentNet(
+    net = TestTimeNet(
         module=DomainAwareModule,
         module__base_module=module,
         module__layer_name=layer_name,
         criterion=TestTimeCriterion,
         criterion__base_criterion=base_criterion,
-        criterion__adapt_criterion=TentLoss(reg_dist, reg_cl, target_criterion),
+        criterion__adapt_criterion=TentLoss(),
         criterion__reg=1,
         **kwargs,
     )

From aab0c09ff04a25ede146c1cf8cb33516e5f271b6 Mon Sep 17 00:00:00 2001
From: mbarneche <aruket.ipa@gmail.com>
Date: Tue, 24 Jun 2025 15:31:56 +0200
Subject: [PATCH 10/19] added Tent, param freeze

---
 skada/deep/_test_time.py                | 85 +++++++++++++++++++++----
 skada/deep/tests/test_deep_test_time.py | 48 +++++++++++---
 2 files changed, 112 insertions(+), 21 deletions(-)

diff --git a/skada/deep/_test_time.py b/skada/deep/_test_time.py
index 2d01c9e1..51fbe10c 100644
--- a/skada/deep/_test_time.py
+++ b/skada/deep/_test_time.py
@@ -63,15 +63,17 @@ def __init__(
         criterion: "TestTimeCriterion",
         optimizer_adapt=None,
         epochs_adapt=None,
+        params_to_adapt=None,
+        layers_to_adapt=None,
         **kwargs,
     ):
         super().__init__(module, criterion=criterion, **kwargs)
         self.optimizer_adapt = optimizer_adapt
         self.epochs_adapt = epochs_adapt
+        self.params_to_adapt = params_to_adapt
+        self.layers_to_adapt = layers_to_adapt
 
-    def fit_source(
-        self, X, y=None, sample_domain=None, sample_weight=None, **fit_params
-    ):
+    def fit(self, X, y=None, sample_domain=None, sample_weight=None, **fit_params):
         print("Training model on source domain...")
         X = self._prepare_input(X, y, sample_domain, sample_weight)
         X = X.select_source()
@@ -86,14 +88,13 @@ def fit_adapt(self, X, sample_domain, sample_weight=None, **fit_params):
         X = self._prepare_input(X, None, sample_domain, sample_weight)
         X = X.select_target()
         self.criterion.train_on_target = True
+        self.freeze_all_params()
+        self.parameters_to_adapt(
+            param_name=self.params_to_adapt, layer_name=self.layers_to_adapt
+        )
         self.partial_fit(X, None, epochs=self.epochs_adapt, **fit_params)
         return self
 
-    def fit(self, X, y=None, sample_domain=None, sample_weight=None, **fit_params):
-        self.fit_source(X, y, sample_domain, sample_weight, **fit_params)
-        self.fit_adapt(X, sample_domain, sample_weight, **fit_params)
-        return self
-
     def initialize_adapt_optimizer(self):
         named_parameters = self.get_all_learnable_params()
         args, kwargs = self.get_params_for_optimizer(
@@ -102,6 +103,69 @@ def initialize_adapt_optimizer(self):
         self.optimizer_ = self.optimizer_adapt(*args, **kwargs)
         return self
 
+    def unfreeze_params(self, module, param_name=None):
+        """
+        Unfreeze the parameters of a given module from the net.
+
+        Parameters
+        ----------
+        param_name : str, list(str), optional
+            The name of the parameter to unfreeze.
+            If None, unfreeze all parameters of the module.
+            If a list, unfreeze all parameters from the list.
+        """
+        # Pseudocode for freezing parameters of a specific layer
+        if param_name is None:
+            param_name = [name for name, _ in module.named_parameters()]
+        elif isinstance(param_name, str):
+            param_name = [param_name]
+
+        for name, param in module.named_parameters():
+            if name in param_name:
+                param.requires_grad = False
+
+        return self
+
+    def parameters_to_adapt(self, param_name=None, layer_name=None):
+        """
+        Unfreeze the parameters of the module.
+        Choose the parameters to adapt during the adaptation phase based on
+        the layer names and parameter names.
+
+        Parameters
+        ----------
+        layer_name : str, list(str), optional
+            The name of the layer to unfreeze.
+            If None, unfreeze params from all layers.
+            If a list, unfreeze all layers in the list.
+        param_name : str, list(str), list(list(str)) optional
+            The name of the parameter to unfreeze.
+            If None, unfreeze all parameters of the layer.
+            If a list, unfreeze all parameters from the list in every layer given.
+            If a list of lists, unfreeze the parameters with the given names in
+            the corresponding layers.
+        """
+        # if None, freeze parameters from all layers of the net
+        if layer_name is None:
+            layer_name = [name for name, _ in self.module.named_modules()]
+        elif isinstance(layer_name, str):
+            layer_name = [layer_name]
+
+        if isinstance(param_name, str):
+            param_name = [param_name] * len(layer_name)
+
+        for module_name, module in self.module.named_modules():
+            if module_name in layer_name:
+                self.unfreeze_params(module=module, param_name=param_name)
+        return self
+
+    def freeze_all_params(self):
+        """Freeze all parameters of the module."""
+        for _, mod in self.module.named_modules():
+            for _, param in mod.named_parameters():
+                param.requires_grad = False
+        return self
+
 
 class TentLoss(BaseDALoss):
     def __init__(self):
@@ -122,9 +186,8 @@ def Tent(
         base_criterion = torch.nn.CrossEntropyLoss()
 
     net = TestTimeNet(
-        module=DomainAwareModule,
-        module__base_module=module,
-        module__layer_name=layer_name,
+        module=DomainAwareModule(module, layer_name),
+        params_to_adapt=["weight", "bias"],
         criterion=TestTimeCriterion,
         criterion__base_criterion=base_criterion,
         criterion__adapt_criterion=TentLoss(),
diff --git a/skada/deep/tests/test_deep_test_time.py b/skada/deep/tests/test_deep_test_time.py
index 1b5c2003..97d7f00e 100644
--- a/skada/deep/tests/test_deep_test_time.py
+++ b/skada/deep/tests/test_deep_test_time.py
@@ -11,21 +11,22 @@
 from torch.utils.data import DataLoader
 
 from skada.datasets import make_shifted_datasets
-from skada.deep._test_time import TestTimeCriterion, TestTimeNet
+from skada.deep._test_time import Tent, TestTimeCriterion, TestTimeNet
 from skada.deep.base import DomainAwareModule
 from skada.deep.losses import TestLoss
 from skada.deep.modules import ToyModule2D
 
 
 @pytest.mark.parametrize(
-    "epochs_adapt, optimizer_adapt",
+    "epochs_adapt, optimizer_adapt, params_to_adapt, layers_to_adapt",
     [
-        (None, None),  # Default test
-        (3, None),  # Test with specific adaptation epochs given
-        (None, Adam),
+        (None, None, None, None),  # Default test
+        (3, None, None, None),  # Test with specific adaptation epochs given
+        (None, Adam, None, None),
+        (None, None, ["weight"], ["dropout"]),
     ],  # Test with Adam optimizer for adaptation
 )
-def test_test_time_criterion(epochs_adapt, optimizer_adapt):
+def test_test_time_net(epochs_adapt, optimizer_adapt, params_to_adapt, layers_to_adapt):
     num_features = 10
     module = ToyModule2D(num_features=num_features)
     criterion = TestTimeCriterion(torch.nn.CrossEntropyLoss(), TestLoss())
@@ -46,6 +47,8 @@ def test_test_time_criterion(epochs_adapt, optimizer_adapt):
         epochs_adapt=epochs_adapt,
         iterator_train=DataLoader,
         criterion=criterion,
+        params_to_adapt=params_to_adapt,
+        layers_to_adapt=layers_to_adapt,
         batch_size=10,
         max_epochs=2,
         train_split=None,
@@ -53,8 +56,33 @@ def test_test_time_criterion(epochs_adapt, optimizer_adapt):
     )
 
     method.fit(X, y, sample_domain)
+    method.fit_adapt(X, y, sample_domain)
 
-    # If the adaptation optimizer is specified, the net's optimizer should have changed
-    # during the fit_adapt method called inside the fit method
-    if optimizer_adapt is not None:
-        assert isinstance(method.optimizer_, optimizer_adapt)
+
+def test_tent():
+    num_features = 10
+    module = ToyModule2D(num_features=num_features)
+
+    n_samples = 20
+    dataset = make_shifted_datasets(
+        n_samples_source=n_samples,
+        n_samples_target=n_samples,
+        shift="conditional_shift",
+        noise=0.1,
+        random_state=42,
+    )
+    X, y, sample_domain = dataset
+    X = X.astype(np.float32)
+
+    method = Tent(
+        DomainAwareModule(module, "dropout"),
+        "dropout",
+        epochs_adapt=3,
+        iterator_train=DataLoader,
+        batch_size=10,
+        max_epochs=2,
+        train_split=None,
+    )
+
+    method.fit(X, y, sample_domain)
+    method.fit_adapt(X, y, sample_domain)

From 616cc29d671203ce04aa02e4bd7b450285bc9fc1 Mon Sep 17 00:00:00 2001
From: mbarneche <aruket.ipa@gmail.com>
Date: Tue, 24 Jun 2025 15:54:29 +0200
Subject: [PATCH 11/19] changed None to all to adapt all params

---
 skada/deep/_test_time.py                | 11 +++++++++--
 skada/deep/tests/test_deep_test_time.py |  1 +
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/skada/deep/_test_time.py b/skada/deep/_test_time.py
index 51fbe10c..6d7cb698 100644
--- a/skada/deep/_test_time.py
+++ b/skada/deep/_test_time.py
@@ -111,11 +111,14 @@ def unfreeze_params(self, module, param_name=None):
         ----------
         param_name : str, list(str), optional
             The name of the parameter to unfreeze.
-            If None, unfreeze all parameters of the module.
+            If 'all', unfreeze all parameters of the module.
+            If None, unfreeze no parameter of the module.
             If a list, unfreeze all parameters from the list.
         """
         # Pseudocode for freezing parameters of a specific layer
         if param_name is None:
+            param_name = []
+        if param_name == "all":
             param_name = [name for name, _ in module.named_parameters()]
         elif isinstance(param_name, str):
             param_name = [param_name]
@@ -136,7 +139,8 @@ def parameters_to_adapt(self, param_name=None, layer_name=None):
         ----------
         layer_name : str, list(str), optional
             The name of the layer to unfreeze.
-            If None, unfreeze params from all layers.
+            If 'all', unfreeze all layers of the module.
+            If None, unfreeze params from no layers.
             If a list, unfreeze all layers in the list.
         param_name : str, list(str), list(list(str)) optional
             The name of the parameter to unfreeze.
@@ -147,6 +151,8 @@ def parameters_to_adapt(self, param_name=None, layer_name=None):
         """
         # if None, freeze parameters from all layers of the net
         if layer_name is None:
+            layer_name = []
+        elif layer_name == "all":
             layer_name = [name for name, _ in self.module.named_modules()]
         elif isinstance(layer_name, str):
             layer_name = [layer_name]
@@ -188,6 +194,7 @@ def Tent(
     net = TestTimeNet(
         module=DomainAwareModule(module, layer_name),
         params_to_adapt=["weight", "bias"],
+        layers_to_adapt="all",
         criterion=TestTimeCriterion,
         criterion__base_criterion=base_criterion,
         criterion__adapt_criterion=TentLoss(),
diff --git a/skada/deep/tests/test_deep_test_time.py b/skada/deep/tests/test_deep_test_time.py
index 97d7f00e..e6c8cb05 100644
--- a/skada/deep/tests/test_deep_test_time.py
+++ b/skada/deep/tests/test_deep_test_time.py
@@ -24,6 +24,7 @@
         (3, None, None, None),  # Test with specific adaptation epochs given
         (None, Adam, None, None),
         (None, None, ["weight"], ["dropout"]),
+        (None, None, "all", "all"),
     ],  # Test with Adam optimizer for adaptation
 )
 def test_test_time_net(epochs_adapt, optimizer_adapt, params_to_adapt, layers_to_adapt):

From b4c1841d715e70fce3c511b24d18b6fa4aaecd9e Mon Sep 17 00:00:00 2001
From: Marion  PAVAUX <marion.pavaux@outlook.fr>
Date: Tue, 24 Jun 2025 15:59:34 +0200
Subject: [PATCH 12/19] Add shot loss

---
 skada/deep/losses.py | 153 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 153 insertions(+)

diff --git a/skada/deep/losses.py b/skada/deep/losses.py
index 21585d7b..30dd1c9f 100644
--- a/skada/deep/losses.py
+++ b/skada/deep/losses.py
@@ -11,12 +11,165 @@
 import skorch  # noqa: F401
 import torch  # noqa: F401
 import torch.nn.functional as F
+from torch.nn import CrossEntropyLoss, LogSoftmax, Module, Softmax
 from torch.nn.functional import mse_loss
 
 from skada.deep.base import BaseDALoss
 from skada.deep.utils import SphericalKMeans
 
 
+class CrossEntropyLabelSmooth(Module):
+    """Estimate the cross-entropy label smooth loss for SHOT method [?]_.
+
+    References
+    ----------
+    .. [?] 	https://doi.org/10.48550/arXiv.2002.08546
+    """
+
+    def __init__(self, num_classes, epsilon=0.1, size_average=True):
+        """Init method"""
+        super().__init__()
+        self.num_classes = num_classes
+        self.epsilon = epsilon
+        self.size_average = size_average
+        self.logsoftmax = LogSoftmax(dim=1)
+
+    def forward(self, y_pred_s, y_s):
+        """Estimate the cross-entropy label smooth loss for SHOT method [?]_.
+
+        Parameters
+        ----------
+        y_s : tensor
+        labels of the source data. Shape (n_samples, n_classes).
+        y_pred_s : tensor
+        predictions of the source data. Shape (n_samples, n_classes).
+
+        Returns
+        -------
+        loss : float
+        The loss of the method.
+        """
+        log_probs = self.logsoftmax(y_pred_s)
+        y_s = torch.zeros(log_probs.size()).scatter_(1, y_s.unsqueeze(1), 1)
+        y_s = (1 - self.epsilon) * y_s + self.epsilon / self.num_classes
+        if self.size_average:
+            loss = (-y_s * log_probs).mean(0).sum()
+        else:
+            loss = (-y_s * log_probs).sum(1)
+        return loss
+
+
+def cross_entropy_label_smooth_loss(
+    y_s: torch.Tensor,
+    y_pred_s: torch.Tensor,
+    num_classes: int,
+    epsilon: float = 0.1,
+    size_average: bool = True,
+) -> float:
+    """Estimate the crpss-entropy label smooth loss loss for SHOT method [?]_.
+
+    Parameters
+    ----------
+    y_s : tensor
+    labels of the source data. Shape (n_samples, n_classes).
+    y_pred_s : tensor
+    predictions of the source data. Shape (n_samples, n_classes).
+    num_classes: int
+    number of classes to predict.
+    epsilon: float
+    smoothing parameter
+    size_average: bool
+
+    Returns
+    -------
+    loss : float
+    The loss of the method.
+    """
+    log_probs = LogSoftmax(dim=1)(y_pred_s)
+    y_s = torch.zeros(log_probs.size()).scatter_(1, y_s.unsqueeze(1), 1)
+    y_s = (1 - epsilon) * y_s + epsilon / num_classes
+    if size_average:
+        loss = (-y_s * log_probs).mean(0).sum()
+    else:
+        loss = (-y_s * log_probs).sum(1)
+    return loss
+
+
+def entropy_loss(y_pred_t_softmax: torch.Tensor) -> float:
+    """Estimate the entropy loss for SHOT method [?]_.
+
+    Parameters
+    ----------
+    y_pred_t_softmax : tensor
+        softmaxed predictions of the target data. Shape (n_samples, n_classes).
+
+    Returns
+    -------
+    loss : float
+        The loss of the method.
+
+    References
+    ----------
+    .. [?] 	https://doi.org/10.48550/arXiv.2002.08546
+    """
+    entropy = -y_pred_t_softmax * torch.log(y_pred_t_softmax + 1e-5)
+    return torch.mean(torch.sum(entropy, axis=1))
+
+
+def diversity_promoting_loss(y_pred_t_softmax: torch.Tensor) -> float:
+    """Estimate the diversity promoting loss for SHOT method [?]_.
+
+    Parameters
+    ----------
+    y_pred_t_softmax : tensor
+        softmaxed predictions of the target data. Shape (n_samples, n_classes).
+
+    Returns
+    -------
+    loss : float
+        The loss of the method.
+
+    References
+    ----------
+    .. [?] 	https://doi.org/10.48550/arXiv.2002.08546
+    """
+    msoftmax = y_pred_t_softmax.mean(dim=0)
+    return torch.sum(msoftmax * torch.log(msoftmax + 1e-5))
+
+
+def shot_full_loss(
+    y_estimate_t: torch.Tensor,
+    y_pred_t: torch.Tensor,
+    entropy_weight: float = 1,
+    div_weight: float = 1,
+    class_weight: float = 0.1,
+) -> float:
+    """Estimate the full loss for SHOT method [?]_.
+
+    Parameters
+    ----------
+    y_estimate_t : tensor
+        estimated labels of the target data, obtained from self-supervised
+        pseudo-labeling strategy. Shape (n_samples, n_classes).
+    y_pred_t : tensor
+        predictions of the target data. Shape (n_samples, n_classes).
+
+    Returns
+    -------
+    loss : float
+        The loss of the method.
+
+    References
+    ----------
+    .. [?] 	https://doi.org/10.48550/arXiv.2002.08546
+    """
+    softmax_y_pred_t = Softmax(dim=1)(y_pred_t)
+    entropy = entropy_loss(softmax_y_pred_t)
+    div = diversity_promoting_loss(softmax_y_pred_t)
+    class_term = CrossEntropyLoss()(y_pred_t, y_estimate_t)
+    return entropy_weight * entropy + div_weight * div + class_term * class_weight
+
+
 def deepcoral_loss(features, features_target, assume_centered=False):
     """Estimate the Frobenius norm divide by 4*n**2
        for DeepCORAL method [12]_.

From 1f52a1be36f2e31414ff764a7b168e1ea2cdbab9 Mon Sep 17 00:00:00 2001
From: Marion  PAVAUX <marion.pavaux@outlook.fr>
Date: Tue, 24 Jun 2025 16:12:06 +0200
Subject: [PATCH 13/19] Add reference for shot loss

---
 README.md            |  2 ++
 skada/deep/losses.py | 20 ++++++++++----------
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index c8cc4dff..0382f2fd 100644
--- a/README.md
+++ b/README.md
@@ -249,3 +249,5 @@ The library is distributed under the 3-Clause BSD license.
 [36] Xiao, Zhiqing, Wang, Haobo, Jin, Ying, Feng, Lei, Chen, Gang, Huang, Fei, Zhao, Junbo.[SPA: A Graph Spectral Alignment Perspective for Domain Adaptation](https://arxiv.org/pdf/2310.17594). In Neurips, 2023.
 
 [37] Xie, Renchunzi, Odonnat, Ambroise, Feofanov, Vasilii, Deng, Weijian, Zhang, Jianfeng and An, Bo. [MaNo: Exploiting Matrix Norm for Unsupervised Accuracy Estimation Under Distribution Shifts](https://arxiv.org/pdf/2405.18979). In NeurIPS, 2024.
+
+[38] Liang, Hu, Feng. [Do We Really Need to Access the Source Data? Source Hypothesis Transfer for Unsupervised Domain Adaptation](https://arxiv.org/abs/2002.08546).
diff --git a/skada/deep/losses.py b/skada/deep/losses.py
index 30dd1c9f..2198c0bd 100644
--- a/skada/deep/losses.py
+++ b/skada/deep/losses.py
@@ -19,11 +19,11 @@
 
 
 class CrossEntropyLabelSmooth(Module):
-    """Estimate the cross-entropy label smooth loss for SHOT method [?]_.
+    """Estimate the cross-entropy label smooth loss for SHOT method [38]_.
 
     References
     ----------
-    .. [?] 	https://doi.org/10.48550/arXiv.2002.08546
+    .. [38] 	https://doi.org/10.48550/arXiv.2002.08546
     """
 
     def __init__(self, num_classes, epsilon=0.1, size_average=True):
@@ -35,7 +35,7 @@ def __init__(self, num_classes, epsilon=0.1, size_average=True):
         self.logsoftmax = LogSoftmax(dim=1)
 
     def forward(self, y_pred_s, y_s):
-        """Estimate the cross-entropy label smooth loss for SHOT method [?]_.
+        """Estimate the cross-entropy label smooth loss for SHOT method [38]_.
 
         Parameters
         ----------
@@ -66,7 +66,7 @@ def cross_entropy_label_smooth_loss(
     epsilon: float = 0.1,
     size_average: bool = True,
 ) -> float:
-    """Estimate the crpss-entropy label smooth loss loss for SHOT method [?]_.
+    """Estimate the crpss-entropy label smooth loss loss for SHOT method [38]_.
 
     Parameters
     ----------
@@ -96,7 +96,7 @@ def cross_entropy_label_smooth_loss(
 
 
 def entropy_loss(y_pred_t_softmax: torch.Tensor) -> float:
-    """Estimate the entropy loss for SHOT method [?]_.
+    """Estimate the entropy loss for SHOT method [38]_.
 
     Parameters
     ----------
@@ -110,14 +110,14 @@ def entropy_loss(y_pred_t_softmax: torch.Tensor) -> float:
 
     References
     ----------
-    .. [?] 	https://doi.org/10.48550/arXiv.2002.08546
+    .. [38] 	https://doi.org/10.48550/arXiv.2002.08546
     """
     entropy = -y_pred_t_softmax * torch.log(y_pred_t_softmax + 1e-5)
     return torch.mean(torch.sum(entropy, axis=1))
 
 
 def diversity_promoting_loss(y_pred_t_softmax: torch.Tensor) -> float:
-    """Estimate the diversity promoting loss for SHOT method [?]_.
+    """Estimate the diversity promoting loss for SHOT method [38]_.
 
     Parameters
     ----------
@@ -131,7 +131,7 @@ def diversity_promoting_loss(y_pred_t_softmax: torch.Tensor) -> float:
 
     References
     ----------
-    .. [?] 	https://doi.org/10.48550/arXiv.2002.08546
+    .. [38] 	https://doi.org/10.48550/arXiv.2002.08546
     """
     msoftmax = y_pred_t_softmax.mean(dim=0)
     return torch.sum(msoftmax * torch.log(msoftmax + 1e-5))
@@ -144,7 +144,7 @@ def shot_full_loss(
     div_weight: float = 1,
     class_weight: float = 0.1,
 ) -> float:
-    """Estimate the full loss for SHOT method [?]_.
+    """Estimate the full loss for SHOT method [38]_.
 
     Parameters
     ----------
@@ -161,7 +161,7 @@ def shot_full_loss(
 
     References
     ----------
-    .. [?] 	https://doi.org/10.48550/arXiv.2002.08546
+    .. [38] 	https://doi.org/10.48550/arXiv.2002.08546
     """
     softmax_y_pred_t = Softmax(dim=1)(y_pred_t)
     entropy = entropy_loss(softmax_y_pred_t)

From bc37c6e12bc3d6535e0a8b6193354068dd3c6b3b Mon Sep 17 00:00:00 2001
From: FDerrida <fr.derrida@gmail.com>
Date: Tue, 24 Jun 2025 16:18:38 +0200
Subject: [PATCH 14/19] Implement SHOT model

---
 skada/deep/modules.py | 91 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 91 insertions(+)

diff --git a/skada/deep/modules.py b/skada/deep/modules.py
index 3563d1a4..7676beb1 100644
--- a/skada/deep/modules.py
+++ b/skada/deep/modules.py
@@ -8,6 +8,7 @@
 import torch
 from torch import nn
 from torch.autograd import Function
+import torch.nn.utils.weight_norm as weightNorm
 
 
 class ToyModule2D(torch.nn.Module):
@@ -226,3 +227,93 @@ def forward(self, x, sample_weight=None):
         x = self.dropout2(x)
         output = self.fc2(x)
         return output
+
+class SHOTNet(nn.Module):
+    """SHOT Network.
+    This network consists of a feature extractor, a bottleneck,
+    and a classifier. The feature extractor is a LeNetBase architecture.
+    """
+
+    def __init__(self, class_num, feature_dim=256, type="ori"):
+        super(SHOTNet, self).__init__()
+        self.feature_extractor = ShotFeatureExtractor()
+        self.bottleneck = FeatBottleneck(
+            feature_dim=self.feature_extractor.in_features,
+            bottleneck_dim=feature_dim,
+            type=type,
+        )
+        self.classifier = FeatClassifier(
+            class_num=class_num, bottleneck_dim=feature_dim, type="linear"
+        )
+
+    def forward(self, x):
+        x = self.feature_extractor(x)
+        x = self.bottleneck(x)
+        x = self.classifier(x)
+        return x
+    
+class ShotFeatureExtractor(nn.Module):
+    """
+    Feature extractor for SHOT.
+    LeNetBase architecture with two convolutional layers,
+    followed by two max pooling layers and a dropout layer.
+    The output is flattened to a vector.
+    """
+    def __init__(self):
+        super(ShotFeatureExtractor, self).__init__()
+        self.conv_params = nn.Sequential(
+                nn.Conv2d(1, 20, kernel_size=5),
+                nn.MaxPool2d(2),
+                nn.ReLU(),
+                nn.Conv2d(20, 50, kernel_size=5),
+                nn.Dropout2d(p=0.5),
+                nn.MaxPool2d(2),
+                nn.ReLU(),
+                )
+        self.in_features = 50*4*4
+
+    def forward(self, x):
+        x = self.conv_params(x)
+        x = x.view(x.size(0), -1)
+        return x
+    
+def init_weights(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv2d') != -1 or classname.find('ConvTranspose2d') != -1:
+        nn.init.kaiming_uniform_(m.weight)
+        nn.init.zeros_(m.bias)
+    elif classname.find('BatchNorm') != -1:
+        nn.init.normal_(m.weight, 1.0, 0.02)
+        nn.init.zeros_(m.bias)
+    elif classname.find('Linear') != -1:
+        nn.init.xavier_normal_(m.weight)
+        nn.init.zeros_(m.bias)
+        
+class FeatBottleneck(nn.Module):
+    def __init__(self, feature_dim, bottleneck_dim=256, type="ori"):
+        super(FeatBottleneck, self).__init__()
+        self.bn = nn.BatchNorm1d(bottleneck_dim, affine=True)
+        self.dropout = nn.Dropout(p=0.5)
+        self.bottleneck = nn.Linear(feature_dim, bottleneck_dim)
+        self.bottleneck.apply(init_weights)
+        self.type = type
+
+    def forward(self, x):
+        x = self.bottleneck(x)
+        if self.type == "bn":
+            x = self.bn(x)
+            x = self.dropout(x)
+        return x
+    
+class FeatClassifier(nn.Module):
+    def __init__(self, class_num, bottleneck_dim=256, type="linear"):
+        super(FeatClassifier, self).__init__()
+        if type == "linear":
+            self.fc = nn.Linear(bottleneck_dim, class_num)
+        else:
+            self.fc = weightNorm(nn.Linear(bottleneck_dim, class_num), name="weight")
+        self.fc.apply(init_weights)
+
+    def forward(self, x):
+        x = self.fc(x)
+        return x

From 6d34ed8a329b0ac5c870bd57632854093e7d800d Mon Sep 17 00:00:00 2001
From: Marion  PAVAUX <marion.pavaux@outlook.fr>
Date: Tue, 24 Jun 2025 18:07:31 +0200
Subject: [PATCH 15/19] Add estimated labels

---
 skada/deep/losses.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/skada/deep/losses.py b/skada/deep/losses.py
index 2198c0bd..785c9f90 100644
--- a/skada/deep/losses.py
+++ b/skada/deep/losses.py
@@ -7,10 +7,12 @@
 
 from functools import partial
 
+import numpy as np
 import ot
 import skorch  # noqa: F401
 import torch  # noqa: F401
 import torch.nn.functional as F
+from scipy.spatial.distance import cdist
 from torch.nn import CrossEntropyLoss, LogSoftmax, Module, Softmax
 from torch.nn.functional import mse_loss
 
@@ -18,6 +20,45 @@
 from skada.deep.utils import SphericalKMeans
 
 
+def get_estimated_label(y_pred_t: torch.Tensor, y_features_t: torch.Tensor) -> int:
+    """Estimate the entropy loss for SHOT method [38]_.
+
+    Parameters
+    ----------
+    y_pred_t : tensor
+        predictions of the target data. Shape (n_samples, n_classes).
+    y_features_t : tensor
+        output of the model just before the classifier. Shape (n_samples, n_classes).
+
+    Returns
+    -------
+    y_estimate_t : int
+        The estimated target prediction.
+
+    References
+    ----------
+    .. [38] 	https://doi.org/10.48550/arXiv.2002.08546
+    """
+    softmax_y_pred_t = Softmax(dim=1)(y_pred_t)
+    n_features = y_pred_t.size(1)
+
+    pred_label = 0
+    for round in range(1):
+        if round == 0:
+            y_estimate_t = (
+                softmax_y_pred_t.float().cpu().numpy()
+                if round == 0
+                else np.eye(n_features)[pred_label]
+            )
+        centroids = y_estimate_t.transpose().dot(y_features_t) / (
+            1e-8 + y_estimate_t.sum(axis=0)[:, None]
+        )
+        cosine_distance = cdist(y_features_t, centroids, "cosine")
+        pred_label = cosine_distance.argmin(axis=1)
+
+    return int(pred_label)
+
+
 class CrossEntropyLabelSmooth(Module):
     """Estimate the cross-entropy label smooth loss for SHOT method [38]_.
 

From 7b1752991850e779a9d6632b15cc8a2fe6d49721 Mon Sep 17 00:00:00 2001
From: Marion  PAVAUX <marion.pavaux@outlook.fr>
Date: Wed, 25 Jun 2025 16:03:18 +0200
Subject: [PATCH 16/19] Add example for test time

---
 examples/deep/plot_test_time.py | 128 ++++++++++++++++++++++++++++++++
 skada/deep/losses.py            |  13 ++--
 skada/deep/modules.py           |  99 +++++++++++++-----------
 3 files changed, 193 insertions(+), 47 deletions(-)
 create mode 100644 examples/deep/plot_test_time.py

diff --git a/examples/deep/plot_test_time.py b/examples/deep/plot_test_time.py
new file mode 100644
index 00000000..5547ba55
--- /dev/null
+++ b/examples/deep/plot_test_time.py
@@ -0,0 +1,128 @@
+"""
+Training setup for deep test time method.
+==========================================
+
+This example illustrates the use of deep test timeS methods in Skada.
+on a simple image classification task.
+"""
+
+# Author: Théo Gnassounou
+#         Marion Pavaux
+#
+# License: BSD 3-Clause
+# sphinx_gallery_thumbnail_number = 4
+
+# %%
+import torch
+from torch.utils.data import DataLoader
+
+from skada.datasets import load_mnist_usps
+from skada.deep.base import DeepDADataset
+from skada.deep.losses import (
+    CrossEntropyLabelSmooth,
+    get_estimated_label,
+    shot_full_loss,
+)
+from skada.deep.modules import SHOTNet
+
+NUM_CLASSES = 2
+
+# %%
+# Load the image datasets
+# ----------------------------------------------------------------------------
+
+sub_dataset = load_mnist_usps(n_classes=NUM_CLASSES, n_samples=0.5)
+dataset = DeepDADataset(*sub_dataset)
+source_dataset = dataset.select_source()
+target_dataset = dataset.select_target()  # TensorDataset(dataset.get_domain("usps"))
+
+# %%
+# Training parameters
+# ----------------------------------------------------------------------------
+
+max_epochs = 100
+batch_size = 256
+lr = 1e-4
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+# %%
+# Training with torch
+# ----------------------------------------------------------------------------
+
+
+model = SHOTNet(class_num=NUM_CLASSES).to(device)
+optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+source_dataloader = DataLoader(source_dataset, batch_size=batch_size, shuffle=True)
+loss_fn = CrossEntropyLabelSmooth(num_classes=NUM_CLASSES)
+
+print("Training loop")
+# Training loop
+for epoch in range(max_epochs):
+    model.train()
+    running_loss = 0.0
+    for batch_idx, (inputs, labels) in enumerate(source_dataloader):
+        inputs, labels = inputs, labels.to(device)
+
+        # Zero the gradients
+        optimizer.zero_grad()
+
+        # Forward pass
+        outputs = model(inputs["X"].to(device))
+        loss = loss_fn(outputs, labels)
+
+        # Backward pass and optimization
+        loss.backward()
+        optimizer.step()
+
+        running_loss += loss.item()
+    print(f"Epoch: {epoch} - Loss: {running_loss / batch_idx}")
+
+target_pred = model(target_dataset.X.to(device)).cpu().detach().numpy().argmax(axis=1)
+accuracy_no_adapation = (target_dataset.y.numpy() == target_pred).mean()
+print(f"Accuracy on target domain without domain adaptation: {accuracy_no_adapation}")
+
+model.classifier.eval()
+optimizer = torch.optim.Adam(
+    [
+        {"params": model.feature_extractor.named_parameters()},
+        {"params": model.bottleneck.named_parameters()},
+    ],
+    lr=lr,
+)
+dataloader_target = DataLoader(target_dataset, batch_size=batch_size, shuffle=True)
+interval_iter = 1
+
+print("Adaptation loop")
+# Adaptation loop
+for epoch in range(max_epochs):
+    model.train()
+    running_loss = 0.0
+    for inputs, labels in dataloader_target:
+        inputs, labels = inputs, labels.to(device)
+
+        # Zero the gradients
+        optimizer.zero_grad()
+
+        if batch_idx % interval_iter == 0:
+            with torch.no_grad():
+                features_full = model.bottleneck(
+                    model.feature_extractor(target_dataset.X.to(device))  # inputs
+                )
+                outputs_full = model.classifier(features_full)
+            estimated_labels = get_estimated_label(outputs_full, features_full)
+
+        # Forward pass
+        outputs = model(inputs["X"].to(device))
+        loss = shot_full_loss(estimated_labels[inputs["sample_idx"]], outputs)
+
+        # Backward pass and optimization
+        loss.backward()
+        optimizer.step()
+
+        running_loss += loss.item()
+    print(f"Epoch: {epoch} - Loss: {running_loss / batch_idx}")
+
+target_pred = model(target_dataset.X.to(device)).cpu().detach().numpy().argmax(axis=1)
+accuracy_no_adapation = (target_dataset.y.numpy() == target_pred).mean()
+print(f"Accuracy on target domain with domain adaptation: {accuracy_no_adapation}")
diff --git a/skada/deep/losses.py b/skada/deep/losses.py
index 785c9f90..02802d70 100644
--- a/skada/deep/losses.py
+++ b/skada/deep/losses.py
@@ -40,6 +40,7 @@ def get_estimated_label(y_pred_t: torch.Tensor, y_features_t: torch.Tensor) -> i
     .. [38] 	https://doi.org/10.48550/arXiv.2002.08546
     """
     softmax_y_pred_t = Softmax(dim=1)(y_pred_t)
+    y_features_t_numpy = y_features_t.cpu().numpy()
     n_features = y_pred_t.size(1)
 
     pred_label = 0
@@ -50,13 +51,13 @@ def get_estimated_label(y_pred_t: torch.Tensor, y_features_t: torch.Tensor) -> i
                 if round == 0
                 else np.eye(n_features)[pred_label]
             )
-        centroids = y_estimate_t.transpose().dot(y_features_t) / (
+        centroids = y_estimate_t.transpose().dot(y_features_t_numpy) / (
             1e-8 + y_estimate_t.sum(axis=0)[:, None]
         )
-        cosine_distance = cdist(y_features_t, centroids, "cosine")
+        cosine_distance = cdist(y_features_t_numpy, centroids, "cosine")
         pred_label = cosine_distance.argmin(axis=1)
 
-    return int(pred_label)
+    return torch.LongTensor(pred_label.astype(int)).to(y_pred_t.device)
 
 
 class CrossEntropyLabelSmooth(Module):
@@ -91,7 +92,9 @@ def forward(self, y_pred_s, y_s):
         The loss of the method.
         """
         log_probs = self.logsoftmax(y_pred_s)
-        y_s = torch.zeros(log_probs.size()).scatter_(1, y_s.unsqueeze(1), 1)
+        y_s = torch.zeros(log_probs.size(), device=y_s.device).scatter_(
+            1, y_s.unsqueeze(1), 1
+        )
         y_s = (1 - self.epsilon) * y_s + self.epsilon / self.num_classes
         if self.size_average:
             loss = (-y_s * log_probs).mean(0).sum()
@@ -175,7 +178,7 @@ def diversity_promoting_loss(y_pred_t_softmax: torch.Tensor) -> float:
     .. [38] 	https://doi.org/10.48550/arXiv.2002.08546
     """
     msoftmax = y_pred_t_softmax.mean(dim=0)
-    return torch.sum(msoftmax * torch.log(msoftmax + 1e-5))
+    return -torch.sum(msoftmax * torch.log(msoftmax + 1e-5))
 
 
 def shot_full_loss(
diff --git a/skada/deep/modules.py b/skada/deep/modules.py
index 7676beb1..e8b81b98 100644
--- a/skada/deep/modules.py
+++ b/skada/deep/modules.py
@@ -6,9 +6,9 @@
 #
 # License: BSD 3-Clause
 import torch
+import torch.nn.utils.weight_norm as weightNorm
 from torch import nn
 from torch.autograd import Function
-import torch.nn.utils.weight_norm as weightNorm
 
 
 class ToyModule2D(torch.nn.Module):
@@ -228,30 +228,7 @@ def forward(self, x, sample_weight=None):
         output = self.fc2(x)
         return output
 
-class SHOTNet(nn.Module):
-    """SHOT Network.
-    This network consists of a feature extractor, a bottleneck,
-    and a classifier. The feature extractor is a LeNetBase architecture.
-    """
-
-    def __init__(self, class_num, feature_dim=256, type="ori"):
-        super(SHOTNet, self).__init__()
-        self.feature_extractor = ShotFeatureExtractor()
-        self.bottleneck = FeatBottleneck(
-            feature_dim=self.feature_extractor.in_features,
-            bottleneck_dim=feature_dim,
-            type=type,
-        )
-        self.classifier = FeatClassifier(
-            class_num=class_num, bottleneck_dim=feature_dim, type="linear"
-        )
 
-    def forward(self, x):
-        x = self.feature_extractor(x)
-        x = self.bottleneck(x)
-        x = self.classifier(x)
-        return x
-    
 class ShotFeatureExtractor(nn.Module):
     """
     Feature extractor for SHOT.
@@ -259,39 +236,46 @@ class ShotFeatureExtractor(nn.Module):
     followed by two max pooling layers and a dropout layer.
     The output is flattened to a vector.
     """
+
     def __init__(self):
-        super(ShotFeatureExtractor, self).__init__()
+        super().__init__()
         self.conv_params = nn.Sequential(
-                nn.Conv2d(1, 20, kernel_size=5),
-                nn.MaxPool2d(2),
-                nn.ReLU(),
-                nn.Conv2d(20, 50, kernel_size=5),
-                nn.Dropout2d(p=0.5),
-                nn.MaxPool2d(2),
-                nn.ReLU(),
-                )
-        self.in_features = 50*4*4
+            nn.Conv2d(1, 20, kernel_size=5),
+            nn.MaxPool2d(2),
+            nn.ReLU(),
+            nn.Conv2d(20, 50, kernel_size=5),
+            nn.Dropout2d(p=0.5),
+            nn.MaxPool2d(2),
+            nn.ReLU(),
+        )
+        self.in_features = 50 * 4 * 4
 
     def forward(self, x):
+        """XXX add docstring here."""
         x = self.conv_params(x)
         x = x.view(x.size(0), -1)
         return x
-    
+
+
 def init_weights(m):
+    """XXX add docstring here."""
     classname = m.__class__.__name__
-    if classname.find('Conv2d') != -1 or classname.find('ConvTranspose2d') != -1:
+    if classname.find("Conv2d") != -1 or classname.find("ConvTranspose2d") != -1:
         nn.init.kaiming_uniform_(m.weight)
         nn.init.zeros_(m.bias)
-    elif classname.find('BatchNorm') != -1:
+    elif classname.find("BatchNorm") != -1:
         nn.init.normal_(m.weight, 1.0, 0.02)
         nn.init.zeros_(m.bias)
-    elif classname.find('Linear') != -1:
+    elif classname.find("Linear") != -1:
         nn.init.xavier_normal_(m.weight)
         nn.init.zeros_(m.bias)
-        
+
+
 class FeatBottleneck(nn.Module):
+    """Feature bottleneck for SHOT."""
+
     def __init__(self, feature_dim, bottleneck_dim=256, type="ori"):
-        super(FeatBottleneck, self).__init__()
+        super().__init__()
         self.bn = nn.BatchNorm1d(bottleneck_dim, affine=True)
         self.dropout = nn.Dropout(p=0.5)
         self.bottleneck = nn.Linear(feature_dim, bottleneck_dim)
@@ -299,15 +283,19 @@ def __init__(self, feature_dim, bottleneck_dim=256, type="ori"):
         self.type = type
 
     def forward(self, x):
+        """XXX add docstring here."""
         x = self.bottleneck(x)
         if self.type == "bn":
             x = self.bn(x)
             x = self.dropout(x)
         return x
-    
+
+
 class FeatClassifier(nn.Module):
+    """Feature classifier for SHOT."""
+
     def __init__(self, class_num, bottleneck_dim=256, type="linear"):
-        super(FeatClassifier, self).__init__()
+        super().__init__()
         if type == "linear":
             self.fc = nn.Linear(bottleneck_dim, class_num)
         else:
@@ -315,5 +303,32 @@ def __init__(self, class_num, bottleneck_dim=256, type="linear"):
         self.fc.apply(init_weights)
 
     def forward(self, x):
+        """XXX add docstring here."""
         x = self.fc(x)
         return x
+
+
+class SHOTNet(nn.Module):
+    """SHOT Network.
+    This network consists of a feature extractor, a bottleneck,
+    and a classifier. The feature extractor is a LeNetBase architecture.
+    """
+
+    def __init__(self, class_num, feature_dim=256, type="ori"):
+        super().__init__()
+        self.feature_extractor = ShotFeatureExtractor()
+        self.bottleneck = FeatBottleneck(
+            feature_dim=self.feature_extractor.in_features,
+            bottleneck_dim=feature_dim,
+            type=type,
+        )
+        self.classifier = FeatClassifier(
+            class_num=class_num, bottleneck_dim=feature_dim, type="linear"
+        )
+
+    def forward(self, x):
+        """XXX add docstring here."""
+        x = self.feature_extractor(x)
+        x = self.bottleneck(x)
+        x = self.classifier(x)
+        return x

From 62d86dbb5e162107dabb75a8acd2fb2e8f69922c Mon Sep 17 00:00:00 2001
From: Marion  PAVAUX <marion.pavaux@outlook.fr>
Date: Wed, 25 Jun 2025 16:25:58 +0200
Subject: [PATCH 17/19] Add example for test time

---
 examples/deep/plot_test_time.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/examples/deep/plot_test_time.py b/examples/deep/plot_test_time.py
index 5547ba55..d42e83f7 100644
--- a/examples/deep/plot_test_time.py
+++ b/examples/deep/plot_test_time.py
@@ -6,8 +6,8 @@
 on a simple image classification task.
 """
 
-# Author: Théo Gnassounou
-#         Marion Pavaux
+# Author: Marion Pavaux
+#         Théo Gnassounou
 #
 # License: BSD 3-Clause
 # sphinx_gallery_thumbnail_number = 4
@@ -34,13 +34,13 @@
 sub_dataset = load_mnist_usps(n_classes=NUM_CLASSES, n_samples=0.5)
 dataset = DeepDADataset(*sub_dataset)
 source_dataset = dataset.select_source()
-target_dataset = dataset.select_target()  # TensorDataset(dataset.get_domain("usps"))
+target_dataset = dataset.select_target()
 
 # %%
 # Training parameters
 # ----------------------------------------------------------------------------
 
-max_epochs = 100
+max_epochs = 10
 batch_size = 256
 lr = 1e-4
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -107,7 +107,7 @@
         if batch_idx % interval_iter == 0:
             with torch.no_grad():
                 features_full = model.bottleneck(
-                    model.feature_extractor(target_dataset.X.to(device))  # inputs
+                    model.feature_extractor(target_dataset.X.to(device))
                 )
                 outputs_full = model.classifier(features_full)
             estimated_labels = get_estimated_label(outputs_full, features_full)

From d29d2995588fd0a432192b05d33f3e45fbfe21cd Mon Sep 17 00:00:00 2001
From: FDerrida <fr.derrida@gmail.com>
Date: Wed, 25 Jun 2025 16:44:26 +0200
Subject: [PATCH 18/19] Add shot implementation and PseudoLabels class

---
 skada/deep/_test_time.py                | 39 +++++++++++-
 skada/deep/callbacks.py                 | 81 +++++++++++++++++++++++++
 skada/deep/losses.py                    | 43 +------------
 skada/deep/modules.py                   | 78 +++++++++++++++++-------
 skada/deep/tests/test_deep_test_time.py | 37 ++++++++++-
 5 files changed, 211 insertions(+), 67 deletions(-)

diff --git a/skada/deep/_test_time.py b/skada/deep/_test_time.py
index 6d7cb698..888509ff 100644
--- a/skada/deep/_test_time.py
+++ b/skada/deep/_test_time.py
@@ -5,7 +5,7 @@
 import torch
 
 from skada.deep.base import BaseDALoss, DomainAwareModule, DomainAwareNet
-from skada.deep.losses import softmax_entropy
+from skada.deep.losses import CrossEntropyLabelSmooth, shot_full_loss, softmax_entropy
 
 
 class TestTimeCriterion(torch.nn.Module):
@@ -81,12 +81,15 @@ def fit(self, X, y=None, sample_domain=None, sample_weight=None, **fit_params):
         self.partial_fit(X, None, **fit_params)
         return self
 
-    def fit_adapt(self, X, sample_domain, sample_weight=None, **fit_params):
+    def fit_adapt(
+        self, X, sample_domain, sample_weight=None, pseudo_labels=None, **fit_params
+    ):
         print("Adapting model to target domain...")
         if self.optimizer_adapt is not None:
             self.initialize_adapt_optimizer()
         X = self._prepare_input(X, None, sample_domain, sample_weight)
         X = X.select_target()
+        X.y = pseudo_labels if pseudo_labels is not None else X.y
         self.criterion.train_on_target = True
         self.freeze_all_params()
         self.parameters_to_adapt(
@@ -203,3 +206,35 @@ def Tent(
     )
 
     return net
+
+
+class ShotLoss(BaseDALoss):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, y_pred, pred_label):
+        loss = shot_full_loss(y_estimate_t=y_pred, y_pred_t=pred_label)
+        return loss
+
+
+def Shot(
+    module,
+    layer_name,
+    base_criterion=CrossEntropyLabelSmooth,
+    **kwargs,
+):
+    if base_criterion is None:
+        base_criterion = torch.nn.CrossEntropyLoss()
+
+    net = TestTimeNet(
+        module=DomainAwareModule(module, layer_name),
+        params_to_adapt=["weight", "bias"],
+        layers_to_adapt="all",
+        criterion=TestTimeCriterion,
+        criterion__base_criterion=base_criterion,
+        criterion__adapt_criterion=ShotLoss(),
+        criterion__reg=1,
+        **kwargs,
+    )
+
+    return net
diff --git a/skada/deep/callbacks.py b/skada/deep/callbacks.py
index 9fc7e6b7..009cc128 100644
--- a/skada/deep/callbacks.py
+++ b/skada/deep/callbacks.py
@@ -4,13 +4,94 @@
 
 import torch
 import torch.nn.functional as F
+from scipy.spatial.distance import cdist
 from skorch.callbacks import Callback
 from skorch.utils import to_tensor
+from torch.nn import Softmax
 
 from skada.deep.base import DomainAwareNet
 from skada.deep.utils import SphericalKMeans
 
 
+class PseudoLabeling(Callback):
+    """Callback to compute pseudo-labels for target domain samples.
+
+    This callback computes pseudo-labels for the target domain samples at the end of
+    each epoch. The pseudo-labels are computed using the memory features and outputs
+    stored in the adaptation criterion.
+
+    """
+
+    def get_estimated_label(y_pred_t: torch.Tensor, y_features_t: torch.Tensor) -> int:
+        """Estimate the entropy loss for SHOT method [38]_.
+
+        Parameters
+        ----------
+        y_pred_t : tensor
+            predictions of the target data. Shape (n_samples, n_classes).
+        y_features_t : tensor
+            output of the model just before the classifier.
+            Shape (n_samples, n_classes).
+
+        Returns
+        -------
+        y_estimate_t : int
+            The estimated target prediction.
+
+        References
+        ----------
+        .. [38] 	https://doi.org/10.48550/arXiv.2002.08546
+        """
+        softmax_y_pred_t = Softmax(dim=1)(y_pred_t)
+        n_features = y_pred_t.size(1)
+
+        pred_label = 0
+        for round in range(1):
+            if round == 0:
+                y_estimate_t = (
+                    softmax_y_pred_t.float().cpu().numpy()
+                    if round == 0
+                    else torch.eye(n_features)[pred_label]
+                )
+            centroids = y_estimate_t.transpose().dot(y_features_t) / (
+                1e-8 + y_estimate_t.sum(axis=0)[:, None]
+            )
+            cosine_distance = cdist(y_features_t, centroids, "cosine")
+            pred_label = cosine_distance.argmin(axis=1)
+
+        return int(pred_label)
+
+    def on_epoch_end(self, net: DomainAwareNet, dataset_train=None, **kwargs):
+        """Compute pseudo-labels at the end of each epoch.
+
+        Parameters
+        ----------
+        net : NeuralNet
+            The neural network being trained.
+        dataset_train : Dataset, optional
+            The training dataset.
+        **kwargs : dict
+            Additional arguments passed to the callback.
+        """
+        X = net._prepare_input(dataset_train)
+
+        # Keep only target samples
+        target = X.select_target()
+        X_t = target.X
+
+        # Disable gradient computation for feature extraction
+        with torch.no_grad():
+            features_t = net.predict_features(X_t)
+            features_t = torch.tensor(features_t, device=net.device)
+
+            # Compute pseudo-labels using get_estimated_label
+            memory_outputs = net.criterion__adapt_criterion.memory_outputs
+            pseudo_labels = self.get_estimated_label(memory_outputs, features_t)
+
+        # Store pseudo-labels in the adaptation criterion
+        net.criterion__adapt_criterion.pseudo_labels = pseudo_labels
+
+
 class ComputeSourceCentroids(Callback):
     """Callback to compute centroids of source domain features for each class.
 
diff --git a/skada/deep/losses.py b/skada/deep/losses.py
index 785c9f90..f6baaf10 100644
--- a/skada/deep/losses.py
+++ b/skada/deep/losses.py
@@ -7,12 +7,12 @@
 
 from functools import partial
 
-import numpy as np
 import ot
 import skorch  # noqa: F401
 import torch  # noqa: F401
 import torch.nn.functional as F
-from scipy.spatial.distance import cdist
+
+# from scipy.spatial.distance import cdist
 from torch.nn import CrossEntropyLoss, LogSoftmax, Module, Softmax
 from torch.nn.functional import mse_loss
 
@@ -20,45 +20,6 @@
 from skada.deep.utils import SphericalKMeans
 
 
-def get_estimated_label(y_pred_t: torch.Tensor, y_features_t: torch.Tensor) -> int:
-    """Estimate the entropy loss for SHOT method [38]_.
-
-    Parameters
-    ----------
-    y_pred_t : tensor
-        predictions of the target data. Shape (n_samples, n_classes).
-    y_features_t : tensor
-        output of the model just before the classifier. Shape (n_samples, n_classes).
-
-    Returns
-    -------
-    y_estimate_t : int
-        The estimated target prediction.
-
-    References
-    ----------
-    .. [38] 	https://doi.org/10.48550/arXiv.2002.08546
-    """
-    softmax_y_pred_t = Softmax(dim=1)(y_pred_t)
-    n_features = y_pred_t.size(1)
-
-    pred_label = 0
-    for round in range(1):
-        if round == 0:
-            y_estimate_t = (
-                softmax_y_pred_t.float().cpu().numpy()
-                if round == 0
-                else np.eye(n_features)[pred_label]
-            )
-        centroids = y_estimate_t.transpose().dot(y_features_t) / (
-            1e-8 + y_estimate_t.sum(axis=0)[:, None]
-        )
-        cosine_distance = cdist(y_features_t, centroids, "cosine")
-        pred_label = cosine_distance.argmin(axis=1)
-
-    return int(pred_label)
-
-
 class CrossEntropyLabelSmooth(Module):
     """Estimate the cross-entropy label smooth loss for SHOT method [38]_.
 
diff --git a/skada/deep/modules.py b/skada/deep/modules.py
index 7676beb1..b269c8d3 100644
--- a/skada/deep/modules.py
+++ b/skada/deep/modules.py
@@ -6,9 +6,9 @@
 #
 # License: BSD 3-Clause
 import torch
+import torch.nn.utils.weight_norm as weightNorm
 from torch import nn
 from torch.autograd import Function
-import torch.nn.utils.weight_norm as weightNorm
 
 
 class ToyModule2D(torch.nn.Module):
@@ -228,14 +228,16 @@ def forward(self, x, sample_weight=None):
         output = self.fc2(x)
         return output
 
+
 class SHOTNet(nn.Module):
-    """SHOT Network.
+    """
+    SHOT Network.
     This network consists of a feature extractor, a bottleneck,
     and a classifier. The feature extractor is a LeNetBase architecture.
     """
 
     def __init__(self, class_num, feature_dim=256, type="ori"):
-        super(SHOTNet, self).__init__()
+        super().__init__()
         self.feature_extractor = ShotFeatureExtractor()
         self.bottleneck = FeatBottleneck(
             feature_dim=self.feature_extractor.in_features,
@@ -247,11 +249,25 @@ def __init__(self, class_num, feature_dim=256, type="ori"):
         )
 
     def forward(self, x):
+        """
+        Forward pass of the SHOTNet.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor after feature extraction, bottleneck, and classification.
+        """
         x = self.feature_extractor(x)
         x = self.bottleneck(x)
         x = self.classifier(x)
         return x
-    
+
+
 class ShotFeatureExtractor(nn.Module):
     """
     Feature extractor for SHOT.
@@ -259,39 +275,52 @@ class ShotFeatureExtractor(nn.Module):
     followed by two max pooling layers and a dropout layer.
     The output is flattened to a vector.
     """
+
     def __init__(self):
-        super(ShotFeatureExtractor, self).__init__()
+        super().__init__()
         self.conv_params = nn.Sequential(
-                nn.Conv2d(1, 20, kernel_size=5),
-                nn.MaxPool2d(2),
-                nn.ReLU(),
-                nn.Conv2d(20, 50, kernel_size=5),
-                nn.Dropout2d(p=0.5),
-                nn.MaxPool2d(2),
-                nn.ReLU(),
-                )
-        self.in_features = 50*4*4
+            nn.Conv2d(1, 20, kernel_size=5),
+            nn.MaxPool2d(2),
+            nn.ReLU(),
+            nn.Conv2d(20, 50, kernel_size=5),
+            nn.Dropout2d(p=0.5),
+            nn.MaxPool2d(2),
+            nn.ReLU(),
+        )
+        self.in_features = 50 * 4 * 4
 
     def forward(self, x):
+        """Forward pass through the bottleneck layer."""
         x = self.conv_params(x)
         x = x.view(x.size(0), -1)
         return x
-    
+
+
 def init_weights(m):
+    """Initialize weights of layers using appropriate initialization schemes.
+
+    Parameters
+    ----------
+    m : nn.Module
+        The layer to initialize.
+    """
     classname = m.__class__.__name__
-    if classname.find('Conv2d') != -1 or classname.find('ConvTranspose2d') != -1:
+    if classname.find("Conv2d") != -1 or classname.find("ConvTranspose2d") != -1:
         nn.init.kaiming_uniform_(m.weight)
         nn.init.zeros_(m.bias)
-    elif classname.find('BatchNorm') != -1:
+    elif classname.find("BatchNorm") != -1:
         nn.init.normal_(m.weight, 1.0, 0.02)
         nn.init.zeros_(m.bias)
-    elif classname.find('Linear') != -1:
+    elif classname.find("Linear") != -1:
         nn.init.xavier_normal_(m.weight)
         nn.init.zeros_(m.bias)
-        
+
+
 class FeatBottleneck(nn.Module):
+    """Bottleneck layer for SHOT architecture."""
+
     def __init__(self, feature_dim, bottleneck_dim=256, type="ori"):
-        super(FeatBottleneck, self).__init__()
+        super().__init__()
         self.bn = nn.BatchNorm1d(bottleneck_dim, affine=True)
         self.dropout = nn.Dropout(p=0.5)
         self.bottleneck = nn.Linear(feature_dim, bottleneck_dim)
@@ -299,15 +328,19 @@ def __init__(self, feature_dim, bottleneck_dim=256, type="ori"):
         self.type = type
 
     def forward(self, x):
+        """Forward pass through the bottleneck layer."""
         x = self.bottleneck(x)
         if self.type == "bn":
             x = self.bn(x)
             x = self.dropout(x)
         return x
-    
+
+
 class FeatClassifier(nn.Module):
+    """Classifier layer for SHOT architecture."""
+
     def __init__(self, class_num, bottleneck_dim=256, type="linear"):
-        super(FeatClassifier, self).__init__()
+        super().__init__()
         if type == "linear":
             self.fc = nn.Linear(bottleneck_dim, class_num)
         else:
@@ -315,5 +348,6 @@ def __init__(self, class_num, bottleneck_dim=256, type="linear"):
         self.fc.apply(init_weights)
 
     def forward(self, x):
+        """Forward pass through the classifier layer."""
         x = self.fc(x)
         return x
diff --git a/skada/deep/tests/test_deep_test_time.py b/skada/deep/tests/test_deep_test_time.py
index e6c8cb05..c881955c 100644
--- a/skada/deep/tests/test_deep_test_time.py
+++ b/skada/deep/tests/test_deep_test_time.py
@@ -11,10 +11,10 @@
 from torch.utils.data import DataLoader
 
 from skada.datasets import make_shifted_datasets
-from skada.deep._test_time import Tent, TestTimeCriterion, TestTimeNet
+from skada.deep._test_time import Shot, Tent, TestTimeCriterion, TestTimeNet
 from skada.deep.base import DomainAwareModule
 from skada.deep.losses import TestLoss
-from skada.deep.modules import ToyModule2D
+from skada.deep.modules import SHOTNet, ToyModule2D
 
 
 @pytest.mark.parametrize(
@@ -87,3 +87,36 @@ def test_tent():
 
     method.fit(X, y, sample_domain)
     method.fit_adapt(X, y, sample_domain)
+
+
+def test_shot():
+    num_features = 10
+    n_classes = 3
+    module = SHOTNet(num_features=num_features)
+
+    n_samples = 20
+    dataset = make_shifted_datasets(
+        n_samples_source=n_samples,
+        n_samples_target=n_samples,
+        shift="conditional_shift",
+        noise=0.1,
+        random_state=42,
+        n_classes=n_classes,
+        n_features=num_features,
+    )
+    X, y, sample_domain = dataset
+    X = X.astype(np.float32)
+
+    method = Shot(
+        DomainAwareModule(module, "dropout"),
+        "dropout",
+        epochs_adapt=3,
+        iterator_train=DataLoader,
+        batch_size=10,
+        max_epochs=2,
+        train_split=None,
+    )
+
+    ## To Do: To modify for this example
+    method.fit(X, y, sample_domain)
+    method.fit_adapt(X, y, sample_domain, pseudo_labels=None)

From d103ef8bc256c08368ea67740e6a4823cfd88691 Mon Sep 17 00:00:00 2001
From: Marion  PAVAUX <marion.pavaux@outlook.fr>
Date: Wed, 25 Jun 2025 17:20:01 +0200
Subject: [PATCH 19/19] Adapt estimated labels to callbacks

---
 examples/deep/plot_test_time.py |  6 ++-
 skada/deep/callbacks.py         | 68 +++++++++------------------------
 skada/deep/losses.py            | 42 --------------------
 skada/deep/utils.py             | 46 +++++++++++++++++++++-
 4 files changed, 66 insertions(+), 96 deletions(-)

diff --git a/examples/deep/plot_test_time.py b/examples/deep/plot_test_time.py
index d42e83f7..602914f2 100644
--- a/examples/deep/plot_test_time.py
+++ b/examples/deep/plot_test_time.py
@@ -20,10 +20,10 @@
 from skada.deep.base import DeepDADataset
 from skada.deep.losses import (
     CrossEntropyLabelSmooth,
-    get_estimated_label,
     shot_full_loss,
 )
 from skada.deep.modules import SHOTNet
+from skada.deep.utils import get_estimated_label_from_centroids
 
 NUM_CLASSES = 2
 
@@ -110,7 +110,9 @@
                     model.feature_extractor(target_dataset.X.to(device))
                 )
                 outputs_full = model.classifier(features_full)
-            estimated_labels = get_estimated_label(outputs_full, features_full)
+            estimated_labels = get_estimated_label_from_centroids(
+                outputs_full, features_full
+            )
 
         # Forward pass
         outputs = model(inputs["X"].to(device))
diff --git a/skada/deep/callbacks.py b/skada/deep/callbacks.py
index 009cc128..f08a86ee 100644
--- a/skada/deep/callbacks.py
+++ b/skada/deep/callbacks.py
@@ -4,13 +4,11 @@
 
 import torch
 import torch.nn.functional as F
-from scipy.spatial.distance import cdist
 from skorch.callbacks import Callback
 from skorch.utils import to_tensor
-from torch.nn import Softmax
 
 from skada.deep.base import DomainAwareNet
-from skada.deep.utils import SphericalKMeans
+from skada.deep.utils import SphericalKMeans, get_estimated_label_from_centroids
 
 
 class PseudoLabeling(Callback):
@@ -22,47 +20,15 @@ class PseudoLabeling(Callback):
 
     """
 
-    def get_estimated_label(y_pred_t: torch.Tensor, y_features_t: torch.Tensor) -> int:
-        """Estimate the entropy loss for SHOT method [38]_.
-
-        Parameters
-        ----------
-        y_pred_t : tensor
-            predictions of the target data. Shape (n_samples, n_classes).
-        y_features_t : tensor
-            output of the model just before the classifier.
-            Shape (n_samples, n_classes).
-
-        Returns
-        -------
-        y_estimate_t : int
-            The estimated target prediction.
-
-        References
-        ----------
-        .. [38] 	https://doi.org/10.48550/arXiv.2002.08546
-        """
-        softmax_y_pred_t = Softmax(dim=1)(y_pred_t)
-        n_features = y_pred_t.size(1)
-
-        pred_label = 0
-        for round in range(1):
-            if round == 0:
-                y_estimate_t = (
-                    softmax_y_pred_t.float().cpu().numpy()
-                    if round == 0
-                    else torch.eye(n_features)[pred_label]
-                )
-            centroids = y_estimate_t.transpose().dot(y_features_t) / (
-                1e-8 + y_estimate_t.sum(axis=0)[:, None]
-            )
-            cosine_distance = cdist(y_features_t, centroids, "cosine")
-            pred_label = cosine_distance.argmin(axis=1)
-
-        return int(pred_label)
-
-    def on_epoch_end(self, net: DomainAwareNet, dataset_train=None, **kwargs):
-        """Compute pseudo-labels at the end of each epoch.
+    def on_epoch_end(
+        self,
+        net: DomainAwareNet,
+        fc_layer: str,
+        dataset_target=None,
+        every_n_epochs: int = 1,
+        **kwargs,
+    ):
+        """TODO Work in progress: Compute pseudo-labels at the end of each epoch.
 
         Parameters
         ----------
@@ -73,20 +39,20 @@ def on_epoch_end(self, net: DomainAwareNet, dataset_train=None, **kwargs):
         **kwargs : dict
             Additional arguments passed to the callback.
         """
-        X = net._prepare_input(dataset_train)
+        X = net._prepare_input(dataset_target)
 
         # Keep only target samples
         target = X.select_target()
         X_t = target.X
 
-        # Disable gradient computation for feature extraction
-        with torch.no_grad():
-            features_t = net.predict_features(X_t)
-            features_t = torch.tensor(features_t, device=net.device)
+        if net.criterion__adapt_criterion.n_epochs % every_n_epochs == 0:
+            with torch.no_grad():
+                features_t = net.predict_features(X_t)
+                features_t = torch.tensor(features_t, device=net.device)
+                outputs_t = eval(f"net.{fc_layer}")(features_t)
 
             # Compute pseudo-labels using get_estimated_label
-            memory_outputs = net.criterion__adapt_criterion.memory_outputs
-            pseudo_labels = self.get_estimated_label(memory_outputs, features_t)
+            pseudo_labels = get_estimated_label_from_centroids(outputs_t, features_t)
 
         # Store pseudo-labels in the adaptation criterion
         net.criterion__adapt_criterion.pseudo_labels = pseudo_labels
diff --git a/skada/deep/losses.py b/skada/deep/losses.py
index 2f39ed60..d628d9f8 100644
--- a/skada/deep/losses.py
+++ b/skada/deep/losses.py
@@ -7,12 +7,10 @@
 
 from functools import partial
 
-import numpy as np
 import ot
 import skorch  # noqa: F401
 import torch  # noqa: F401
 import torch.nn.functional as F
-from scipy.spatial.distance import cdist
 
 # from scipy.spatial.distance import cdist
 from torch.nn import CrossEntropyLoss, LogSoftmax, Module, Softmax
@@ -22,46 +20,6 @@
 from skada.deep.utils import SphericalKMeans
 
 
-def get_estimated_label(y_pred_t: torch.Tensor, y_features_t: torch.Tensor) -> int:
-    """Estimate the entropy loss for SHOT method [38]_.
-
-    Parameters
-    ----------
-    y_pred_t : tensor
-        predictions of the target data. Shape (n_samples, n_classes).
-    y_features_t : tensor
-        output of the model just before the classifier. Shape (n_samples, n_classes).
-
-    Returns
-    -------
-    y_estimate_t : int
-        The estimated target prediction.
-
-    References
-    ----------
-    .. [38] 	https://doi.org/10.48550/arXiv.2002.08546
-    """
-    softmax_y_pred_t = Softmax(dim=1)(y_pred_t)
-    y_features_t_numpy = y_features_t.cpu().numpy()
-    n_features = y_pred_t.size(1)
-
-    pred_label = 0
-    for round in range(1):
-        if round == 0:
-            y_estimate_t = (
-                softmax_y_pred_t.float().cpu().numpy()
-                if round == 0
-                else np.eye(n_features)[pred_label]
-            )
-        centroids = y_estimate_t.transpose().dot(y_features_t_numpy) / (
-            1e-8 + y_estimate_t.sum(axis=0)[:, None]
-        )
-        cosine_distance = cdist(y_features_t_numpy, centroids, "cosine")
-        pred_label = cosine_distance.argmin(axis=1)
-
-    return torch.LongTensor(pred_label.astype(int)).to(y_pred_t.device)
-
-
 class CrossEntropyLabelSmooth(Module):
     """Estimate the cross-entropy label smooth loss for SHOT method [38]_.
 
diff --git a/skada/deep/utils.py b/skada/deep/utils.py
index 40fed402..2302c5b5 100644
--- a/skada/deep/utils.py
+++ b/skada/deep/utils.py
@@ -6,11 +6,13 @@
 import numbers
 from functools import partial
 
+import numpy as np
 from skorch.utils import _identity
 import torch
-from torch.nn import CrossEntropyLoss
+from torch.nn import CrossEntropyLoss, Softmax
 from torch.nn.functional import cosine_similarity
 from sklearn.utils.validation import check_is_fitted
+from scipy.spatial.distance import cdist
 
 
 def _get_intermediate_layers(intermediate_layers, layer_name):
@@ -333,3 +335,45 @@ def _compute_dissimilarity_loss(self, X, centroids):
         dissimilarity_loss = dissimilarities.sum().item()
 
         return dissimilarity_loss
+
+
+def get_estimated_label_from_centroids(
+    y_pred_t: torch.Tensor, y_features_t: torch.Tensor
+) -> int:
+    """Estimate the entropy loss for SHOT method [38]_.
+
+    Parameters
+    ----------
+    y_pred_t : tensor
+        predictions of the target data. Shape (n_samples, n_classes).
+    y_features_t : tensor
+        output of the model just before the classifier. Shape (n_samples, n_classes).
+
+    Returns
+    -------
+    y_estimate_t : int
+        The estimated target prediction.
+
+    References
+    ----------
+    .. [38] 	https://doi.org/10.48550/arXiv.2002.08546
+    """
+    softmax_y_pred_t = Softmax(dim=1)(y_pred_t)
+    y_features_t_numpy = y_features_t.cpu().numpy()
+    n_features = y_pred_t.size(1)
+
+    pred_label = 0
+    for round in range(1):
+        if round == 0:
+            y_estimate_t = (
+                softmax_y_pred_t.float().cpu().numpy()
+                if round == 0
+                else np.eye(n_features)[pred_label]
+            )
+        centroids = y_estimate_t.transpose().dot(y_features_t_numpy) / (
+            1e-8 + y_estimate_t.sum(axis=0)[:, None]
+        )
+        cosine_distance = cdist(y_features_t_numpy, centroids, "cosine")
+        pred_label = cosine_distance.argmin(axis=1)
+
+    return torch.LongTensor(pred_label.astype(int)).to(y_pred_t.device)