Nit

tensorflow · Nov 4, 2022 · de84a0e · de84a0e
1 parent 8976a0f
commit de84a0e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 8 deletions.
diff --git a/e2e/benchmarks/local-benchmark/index.html b/e2e/benchmarks/local-benchmark/index.html
@@ -174,17 +174,14 @@ <h2>TensorFlow.js Model Benchmark</h2>
     async function predictAndGetData(predict, model, inferenceInput, enableDump) {
       const prediction = await predict(model, inferenceInput);
       let intermediateData = {};
-      let graphModel = null;
       if (enableDump) {
-        graphModel = getGraphModel(model);
+        const graphModel = getGraphModel(model);
         if (graphModel) {
           intermediateData = await getIntermediateTensorInfo(graphModel.getIntermediateTensors());
+          graphModel.disposeIntermediateTensors();
         }
       }
       const predictionData = await getPredictionData(prediction);
-      if (graphModel) {
-        graphModel.disposeIntermediateTensors();
-      }
       return {data: predictionData, intermediateData};
     }
 

diff --git a/tfjs-converter/src/executor/graph_executor.ts b/tfjs-converter/src/executor/graph_executor.ts
@@ -272,13 +272,13 @@ export class GraphExecutor implements FunctionExecutor {
               outputNodeNames, intermediateTensorConsumerCount);
         }
       }
+      if (this.keepIntermediateTensors) {
+        this.clonedTensorsMap = this.cloneTensorMap(tensorsMap);
+      }
       // dispose the context for the root executor
       if (this.parent == null) {
         context.dispose(tensorsToKeep);
       }
-      if (this.keepIntermediateTensors) {
-        this.clonedTensorsMap = this.cloneTensorMap(tensorsMap);
-      }
 
       return outputs.map(name => getTensor(name, tensorsMap, context));
     });