Fixed a small bug in metric wrapper

Author: Seilmast

CollaborativeCoding/load_metric.py

@@ -79,6 +79,7 @@ def _get_metric(self, key):
                 raise ValueError(f"Metric {key} not supported")
 
     def __call__(self, y_true, y_pred):
+        y_true, y_pred = y_true.detach().cpu(), y_pred.detach().cpu()
         for key in self.metrics:
             self.metrics[key](y_true, y_pred)
 

CollaborativeCoding/metrics/EntropyPred.py

@@ -36,7 +36,6 @@ def __call__(self, y_true: th.Tensor, y_logits: th.Tensor):
         assert y_logits.size(-1) == self.num_classes, (
             f"y_logit class length: {y_logits.size(-1)}, expected: {self.num_classes}"
         )
-
         y_pred = nn.Softmax(dim=1)(y_logits)
         print(f"y_pred: {y_pred}")
         entropy_values = entropy(y_pred, axis=1)