fix batch sizes

tensorflow · Dec 11, 2020 · c79ac9b · c79ac9b
1 parent 1632a0b
commit c79ac9b
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/TrainingLoop/Metrics.swift b/TrainingLoop/Metrics.swift
@@ -173,9 +173,10 @@ public struct AccuracyMeasurerTop5: MetricsMeasurer {
         "For accuracy measurements, the model output must be Tensor<Float>, and the labels must be Tensor<Int>."
       )
     }
-    let top5 = _Raw.topKV2(predictions, k: Tensor<Int32>(5), sorted:false)
-    let labelsReshaped = labels.reshaped(to: [1, 32])
-    let labelsTranspose = labelsReshaped.transposed(permutation: [1,0])
+    let top5 = _Raw.topKV2(predictions, k: Tensor<Int32>(5), sorted: false)
+    let batchSize = labels.shape[0]
+    let labelsReshaped = labels.reshaped(to: [1, batchSize])
+    let labelsTranspose = labelsReshaped.transposed(permutation: [1, 0])
     let ones = Tensor<Int32>([1, 1, 1, 1, 1])
     let exandedLabels = labelsTranspose * ones
     correctGuessCount += Tensor<Int32>(top5.indices .== exandedLabels).sum().scalarized()