Minor refactoring for better comparison

JetBrains-Research · Aug 22, 2024 · 60bbde0 · 60bbde0
1 parent 7d46963
commit 60bbde0
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/inference/inference-core/src/commonMain/kotlin/io/kinference.core/model/KIModel.kt b/inference/inference-core/src/commonMain/kotlin/io/kinference.core/model/KIModel.kt
@@ -42,10 +42,10 @@ class KIModel(val id: String, val name: String, val opSet: OperatorSetRegistry,
         val arrayContext = coroutineContext[ArrayContext.Key]
         val results = if (arrayContext == null) {
             withContext(ArrayContext()) {
-                graph.execute(input, contexts).also { it.forEach { it.markOutput() } }
+                graph.execute(input, contexts).onEach { it.markOutput() }
             }
         } else {
-            graph.execute(input, contexts).also { it.forEach { it.markOutput() } }
+            graph.execute(input, contexts).onEach { it.markOutput() }
         }
 //        val results = withContext(ArrayContext()) {
 //            graph.execute(input, contexts).also { it.forEach { it.markOutput() } }

diff --git a/utils/utils-testing/src/commonMain/kotlin/io.kinference/runners/PerformanceRunner.kt b/utils/utils-testing/src/commonMain/kotlin/io.kinference/runners/PerformanceRunner.kt
@@ -54,7 +54,7 @@ class PerformanceRunner<T : ONNXData<*, *>>(private val engine: TestEngine<T>) {
         val results = ArrayList<PerformanceResults>()
 
         withContext(ArrayContext()) {
-            val arrayContext = coroutineContext[ArrayContext.Key]
+//            val arrayContext = coroutineContext[ArrayContext.Key]
             for (dataset in datasets) {
                 val inputs = dataset.data.map { engine.loadData(it.first, it.second) }
 
@@ -90,7 +90,8 @@ class PerformanceRunner<T : ONNXData<*, *>>(private val engine: TestEngine<T>) {
                         }.millis
                         times[i] = time
 
-                        outputs.values.forEach { arrayContext?.returnNDArray(it.data as NDArrayCore) }
+                        outputs.values.forEach { it.close() }
+//                        outputs.values.forEach { arrayContext?.returnNDArray(it.data as NDArrayCore) }
                     }
                 }