Skip to content

Commit

Permalink
fix device of macro-f1 in evaluation
Browse files Browse the repository at this point in the history
  • Loading branch information
sfluegel committed Jan 4, 2024
1 parent c131481 commit bddd271
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion chebai/result/classification.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ def print_metrics(preds, labels, device, classes=None, top_k=10, markdown_output
f1_macro = MultilabelF1Score(preds.shape[1], average="macro").to(device=device)
f1_micro = MultilabelF1Score(preds.shape[1], average="micro").to(device=device)

my_f1_macro = MacroF1(preds.shape[1])
my_f1_macro = MacroF1(preds.shape[1]).to(device=device)

print(
f"Macro-F1 (torchmetrics, unadjusted) on test set with {preds.shape[1]} classes: {f1_macro(preds, labels):3f}"
Expand Down

0 comments on commit bddd271

Please sign in to comment.