cast to cpu before numpy()

avoids error if trying to call .numpy() on tensor that is not on cpu
kitzeslab · Nov 11, 2024 · b8161d4 · b8161d4
1 parent 028051a
commit b8161d4
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/opensoundscape/ml/shallow_classifier.py b/opensoundscape/ml/shallow_classifier.py
@@ -141,13 +141,15 @@ def quick_fit(
                 )
                 try:
                     auroc = roc_auc_score(
-                        validation_labels.detach().numpy(), val_outputs.detach().numpy()
+                        validation_labels.detach().cpu().numpy(),
+                        val_outputs.detach().cpu().numpy(),
                     )
                 except:
                     auroc = float("nan")
                 try:
                     map = average_precision_score(
-                        validation_labels.detach().numpy(), val_outputs.detach().numpy()
+                        validation_labels.detach().cpu().numpy(),
+                        val_outputs.detach().cpu().numpy(),
                     )
                 except:
                     map = float("nan")