fix cpu/cuda mismatch

EleutherAI · Apr 24, 2023 · 6fb6ef7 · 6fb6ef7
1 parent ad4cf34
commit 6fb6ef7
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/elk/training/supervised.py b/elk/training/supervised.py
@@ -34,7 +34,11 @@ def evaluate_supervised(
         raise ValueError(f"Invalid val_h shape: {val_h.shape}")
 
     lr_acc = accuracy(labels, raw_preds.flatten())
-    lr_auroc = RocAucResult(-1., -1., -1.) if len(labels.unique()) == 1 else roc_auc_ci(labels, logits.flatten())
+    lr_auroc = (
+        RocAucResult(-1.0, -1.0, -1.0)
+        if len(labels.unique()) == 1
+        else roc_auc_ci(labels.cpu(), logits.cpu().flatten())
+    )
 
     return lr_auroc, assert_type(float, lr_acc)