fixed metrics and weights

models/utils.py

@@ -37,7 +37,9 @@ def calculate_metrics(
     pred, target, threshold=0.5, prefix="", multi_label=True
 ) -> dict[str, torch.Tensor]:
     target = target.detach().cpu().numpy()
-    pred = pred.detach().cpu().numpy()
+    pred = pred.detach().cpu()
+    pred = nn.functional.softmax(pred, dim=1)
+    pred = pred.numpy()
     params = {
         "y_true": target if multi_label else target.argmax(1),
         "y_pred": np.array(pred > threshold, dtype=float)

preprocessing/dataset.py

@@ -80,7 +80,9 @@ class SongDataset(Dataset):
 
     def get_label_weights(self):
         n_examples, n_classes = self.dance_labels.shape
-        return torch.from_numpy(n_examples / (n_classes * sum(self.dance_labels)))
+        weights = n_examples / (n_classes * sum(self.dance_labels))
+        weights[np.isinf(weights)] = 0.0
+        return torch.from_numpy(weights)
 
     def _backtrace_audio_path(self, index: int) -> str:
         return self.audio_paths[self._idx2audio_idx(index)]