Fix a bug

Author: lizekai-richard

ddranking/metrics/soft_label.py

@@ -148,7 +148,7 @@ def load_real_data(self, dataset, class_map, num_classes):
 
     def get_class_indices(self, dataset, class_map, num_classes):
         class_indices = [[] for c in range(num_classes)]
-        for idx, (_, label) in enumerate(dataset.imgs):
+        for idx, label in enumerate(dataset.targets):
             if torch.is_tensor(label):
                 label = label.item()
             true_label = class_map[label]

ddranking/utils/data.py

@@ -38,12 +38,14 @@ class Config:
 class TensorDataset(torch.utils.data.Dataset):
 
     def __init__(self, images: Tensor, labels: Tensor):
+        super(TensorDataset, self).__init__()
         self.images = images
-        self.labels = labels
+        self.targets = labels
+        self.imgs = [(image, label) for image, label in zip(images, labels)]
 
     def __getitem__(self, index: int):
         image = self.images[index]
-        label = self.labels[index]
+        label = self.targets[index]
         return image, label
 
     def __len__(self):