Repository: pykale/pykale
Issue: [Bug] Getting indicies with `torch.where` may return indicies of multiple tensor modes.

# 🐛 Bug
In `kale.pipeline.multi_domain_adapter`, torch.where may return indicies of both rows and columns when the shape of `domain_labels` is (num_samples, 1), which as a result only the first column of data features will be taken.



Changed files:
tests/loaddata/test_image_access.py, kale/pipeline/multi_domain_adapter.py, kale/loaddata/multi_domain.py, kale/predict/losses.py

diff --git a/kale/loaddata/multi_domain.py b/kale/loaddata/multi_domain.py
index 0eb9011..d93368c 100644
--- a/kale/loaddata/multi_domain.py
+++ b/kale/loaddata/multi_domain.py
@@ -536,10 +536,11 @@ class MultiDomainAdapDataset(DomainsDatasetBase):
         val_split_ratio (float, optional): Split ratio for validation set. Defaults to 0.1.
         test_split_ratio (float, optional): Split ratio for test set. Defaults to 0.2.
         random_state (int, optional): Random state for generator. Defaults to 1.
+        test_on_all (bool, optional): Whether test model on all target. Defaults to False.
     """
 
     def __init__(
-        self, data_access, val_split_ratio=0.1, test_split_ratio=0.2, random_state: int = 1,
+        self, data_access, val_split_ratio=0.1, test_split_ratio=0.2, random_state: int = 1, test_on_all=False
     ):
         self.domain_to_idx = data_access.domain_to_idx
         self.n_domains = len(data_access.domain_to_idx)
@@ -550,17 +551,24 @@ class MultiDomainAdapDataset(DomainsDatasetBase):
         self._sampling_config = FixedSeedSamplingConfig(seed=random_state, balance_domain=True)
         self._loader = MultiDataLoader
         self._random_state = random_state
+        self.test_on_all = test_on_all
 
     def prepare_data_loaders(self):
         splits = ["test", "valid", "train"]
         self._sample_by_split["test"] = self.data_access.get_test()
         if self._sample_by_split["test"] is None:
             # split test, valid, and train set if the data access no train test splits
-            subset_idx = _domain_stratified_split(
-                self.data_access.domain_labels, 3, [self._test_split_ratio, self._val_split_ratio]
-            )
-            for i in range(len(splits)):
-                self._sample_by_split[splits[i]] = torch.utils.data.Subset(self.data_access, subset_idx[i])
+            if self.test_on_all:
+                subset_idx = _domain_stratified_split(self.data_access.domain_labels, 2, [self._val_split_ratio])
+                self._sample_by_split["valid"] = torch.utils.data.Subset(self.data_access, subset_idx[0])
+                self._sample_by_split["train"] = torch.utils.data.Subset(self.data_access, subset_idx[1])
+                self._sample_by_split["test"] = torch.utils.data.Subset(self.data_access, np.concatenate(subset_idx))
+            else:
+                subset_idx = _domain_stratified_split(
+                    self.data_access.domain_labels, 3, [self._test_split_ratio, self._val_split_ratio]
+                )
+                for i in range(len(splits)):
+                    self._sample_by_split[splits[i]] = torch.utils.data.Subset(self.data_access, subset_idx[i])
         else:
             # use original data split if get_test() is not none
             self._sample_by_split["valid"], self._sample_by_split["train"] = self.data_access.get_train_val(
diff --git a/kale/pipeline/multi_domain_adapter.py b/kale/pipeline/multi_domain_adapter.py
index c10ad4c..1f5226b 100644
--- a/kale/pipeline/multi_domain_adapter.py
+++ b/kale/pipeline/multi_domain_adapter.py
@@ -128,8 +128,8 @@ class M3SDATrainer(BaseMultiSourceTrainer):
         x, y, domain_labels = batch
         phi_x = self.forward(x)
         moment_loss = self._compute_domain_dist(phi_x, domain_labels)
-        src_idx = torch.where(domain_labels != self.target_label)
-        tgt_idx = torch.where(domain_labels == self.target_label)
+        src_idx = torch.where(domain_labels != self.target_label)[0]
+        tgt_idx = torch.where(domain_labels == self.target_label)[0]
         cls_loss, ok_src = self._compute_cls_loss(phi_x[src_idx], y[src_idx], domain_labels[src_idx])
         if len(tgt_idx) > 0:
             y_tgt_hat = _average_cls_output(phi_x[tgt_idx], self.classifiers)
@@ -156,7 +156,7 @@ class M3SDATrainer(BaseMultiSourceTrainer):
             for domain_ in self.domain_to_idx.keys():
                 if domain_ == self.target_domain:
                     continue
-                domain_idx = torch.where(domain_labels == self.domain_to_idx[domain_])
+                domain_idx = torch.where(domain_labels == self.domain_to_idx[domain_])[0]
                 cls_output = self.classifiers[domain_](x[domain_idx])
                 loss_cls_, ok_src_ = losses.cross_entropy_logits(cls_output, y[domain_idx])
                 cls_loss += loss_cls_
@@ -211,8 +211,8 @@ class _DINTrainer(BaseMultiSourceTrainer):
         x, y, domain_labels = batch
         phi_x = self.forward(x)
         loss_dist = self._compute_domain_dist(phi_x, domain_labels)
-        src_idx = torch.where(domain_labels != self.target_label)
-        tgt_idx = torch.where(domain_labels == self.target_label)
+        src_idx = torch.where(domain_labels != self.target_label)[0]
+        tgt_idx = torch.where(domain_labels == self.target_label)[0]
         cls_output = self.classifier(phi_x)
         loss_cls, ok_src = losses.cross_entropy_logits(cls_output[src_idx], y[src_idx])
         _, ok_tgt = losses.cross_entropy_logits(cls_output[tgt_idx], y[tgt_idx])
@@ -275,13 +275,13 @@ class MFSANTrainer(BaseMultiSourceTrainer):
     def compute_loss(self, batch, split_name="V"):
         x, y, domain_labels = batch
         phi_x = self.forward(x)
-        tgt_idx = torch.where(domain_labels == self.target_label)
+        tgt_idx = torch.where(domain_labels == self.target_label)[0]
         n_src = len(self.src_domains)
         domain_dist = 0
         loss_cls = 0
         ok_src = []
         for src_domain in self.src_domains:
-            src_domain_idx = torch.where(domain_labels == self.domain_to_idx[src_domain])
+            src_domain_idx = torch.where(domain_labels == self.domain_to_idx[src_domain])[0]
             phi_src = self.domain_net[src_domain].forward(phi_x[src_domain_idx])
             phi_tgt = self.domain_net[src_domain].forward(phi_x[tgt_idx])
             kernels = losses.gaussian_kernel(
diff --git a/kale/predict/losses.py b/kale/predict/losses.py
index 9a251d3..cebc903 100644
--- a/kale/predict/losses.py
+++ b/kale/predict/losses.py
@@ -168,7 +168,7 @@ def _moment_k(x: torch.Tensor, domain_labels: torch.Tensor, k_order=2):
     n_unique_domain_ = len(unique_domain_)
     x_k_order = []
     for domain_label_ in unique_domain_:
-        domain_idx = torch.where(domain_labels == domain_label_)
+        domain_idx = torch.where(domain_labels == domain_label_)[0]
         x_mean = x[domain_idx].mean(0)
         if k_order == 1:
             x_k_order.append(x_mean)
diff --git a/tests/loaddata/test_image_access.py b/tests/loaddata/test_image_access.py
index 6277faf..8168d89 100644
--- a/tests/loaddata/test_image_access.py
+++ b/tests/loaddata/test_image_access.py
@@ -8,13 +8,14 @@ from kale.loaddata.image_access import DigitDataset, DigitDatasetAccess, get_cif
 from kale.loaddata.multi_domain import DomainsDatasetBase, MultiDomainAdapDataset, MultiDomainDatasets
 
 
-def test_office31(office_path):
+@pytest.mark.parametrize("test_on_all", [True, False])
+def test_office31(office_path, test_on_all):
     office_access = ImageAccess.get_multi_domain_images(
         "OFFICE31", office_path, download=True, return_domain_label=True
     )
     testing.assert_equal(len(office_access.class_to_idx), 31)
     testing.assert_equal(len(office_access.domain_to_idx), 3)
-    dataset = MultiDomainAdapDataset(office_access)
+    dataset = MultiDomainAdapDataset(office_access, test_on_all=test_on_all)
     dataset.prepare_data_loaders()
     domain_labels = list(dataset.domain_to_idx.values())
     for split in ["train", "valid", "test"]: