Use memory statistics from model probing in calculation

1 files changed, 3 insertions, 2 deletions

diff --git a/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxModelCost.java b/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxModelCost.java
index 422ceba8074..595cd97e6b6 100644
--- a/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxModelCost.java
+++ b/config-model-api/src/main/java/com/yahoo/config/model/api/OnnxModelCost.java
@@ -4,6 +4,7 @@ package com.yahoo.config.model.api;
 
 import com.yahoo.config.ModelReference;
 import com.yahoo.config.application.api.ApplicationFile;
+import com.yahoo.config.application.api.ApplicationPackage;
 import com.yahoo.config.application.api.DeployLogger;
 
 /**
@@ -11,7 +12,7 @@ import com.yahoo.config.application.api.DeployLogger;
  */
 public interface OnnxModelCost {
 
-    Calculator newCalculator(DeployLogger logger);
+    Calculator newCalculator(ApplicationPackage appPkg, DeployLogger logger);
 
     interface Calculator {
         long aggregatedModelCostInBytes();
@@ -20,7 +21,7 @@ public interface OnnxModelCost {
     }
 
     static OnnxModelCost disabled() {
-        return (__) -> new Calculator() {
+        return (__, ___) -> new Calculator() {
             @Override public long aggregatedModelCostInBytes() { return 0; }
             @Override public void registerModel(ApplicationFile path) {}
             @Override public void registerModel(ModelReference ref) {}