fixup! discojs/validation: argmax for multiclass preds

epfml · Nov 21, 2022 · 8e3c573 · 8e3c573
1 parent ee982e6
commit 8e3c573
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/discojs/discojs-core/src/validation/validator.ts b/discojs/discojs-core/src/validation/validator.ts
@@ -45,12 +45,8 @@ export class Validator {
         this.size += xs.shape[0]
 
         // Get labels from one hot encoding
-        const preds = List(oneHotPreds.reshape([-1, this.classes]).arraySync() as number[][])
-          .map(List)
-          .map((ps) => ps.indexOf(ps.max() ?? 0))
-        const labels = List(ys.reshape([-1, this.classes]).arraySync() as number[][])
-          .map(List)
-          .map((ps) => ps.indexOf(1))
+        const preds = List(oneHotPreds.reshape([-1, this.classes]).argMax(1).arraySync() as number[])
+        const labels = List(ys.reshape([-1, this.classes]).argMax(1).arraySync() as number[])
 
         // Keep track of prediction results for the confusion matrix
         this.preds = this.preds.push(...preds)