[TO_REVIEW] Fix deep coral

skada/deep/losses.py

@@ -1,6 +1,7 @@
 # Author: Theo Gnassounou <[email protected]>
 #         Remi Flamary <[email protected]>
 #         Yanis Lalou <[email protected]>
+#         Antoine Collas <[email protected]>
 #
 # License: BSD 3-Clause
 
@@ -9,6 +10,7 @@
 import ot
 import skorch  # noqa: F401
 import torch  # noqa: F401
+from torch.nn.functional import mse_loss
 
 from skada.deep.base import BaseDALoss
 
@@ -40,10 +42,11 @@ def deepcoral_loss(features, features_target, assume_centered=False):
     if not assume_centered:
         features = features - features.mean(0)
         features_target = features_target - features_target.mean(0)
-    cov = torch.cov(features)
-    cov_target = torch.cov(features_target)
-    diff = cov - cov_target
-    loss = (diff * diff).sum() / (4 * len(cov) ** 2)
+    cov = torch.cov(features.T)
+    cov_target = torch.cov(features_target.T)
+    divergence = mse_loss(cov, cov_target, reduction="sum")
+    dim = features.shape[1]
+    loss = (1 / (4 * (dim**2))) * divergence
     return loss
 
 

skada/deep/tests/test_deep_divergence.py

@@ -39,7 +39,7 @@ def test_deepcoral(assume_centered):
         ToyModule2D(),
         reg=1,
         layer_name="dropout",
-        batch_size=10,
+        batch_size=15,
         max_epochs=10,
         train_split=None,
         assume_centered=assume_centered,