From ee80a7f404fedf6fd84baabc74a8605943a72005 Mon Sep 17 00:00:00 2001
From: Miguel Alonso Jr <76960110+miguelalonsojr@users.noreply.github.com>
Date: Thu, 30 Nov 2023 07:20:16 -0500
Subject: [PATCH] Fix tensor disposal (#6028)

* Upgraded to Sentis 1.3.0-pre.2

* Fixed tensor disposal bug in ModelRunner.
---
 Project/Packages/packages-lock.json                  | 4 ++--
 com.unity.ml-agents/CHANGELOG.md                     | 1 +
 com.unity.ml-agents/Runtime/Inference/ModelRunner.cs | 4 ++++
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/Project/Packages/packages-lock.json b/Project/Packages/packages-lock.json
index 2bfe2876bd..8066828231 100644
--- a/Project/Packages/packages-lock.json
+++ b/Project/Packages/packages-lock.json
@@ -72,7 +72,7 @@
       "depth": 0,
       "source": "local",
       "dependencies": {
-        "com.unity.sentis": "1.3.0-pre.1",
+        "com.unity.sentis": "1.3.0-pre.2",
         "com.unity.modules.imageconversion": "1.0.0",
         "com.unity.modules.jsonserialize": "1.0.0",
         "com.unity.modules.physics": "1.0.0"
@@ -104,7 +104,7 @@
       "url": "https://packages.unity.com"
     },
     "com.unity.sentis": {
-      "version": "1.3.0-pre.1",
+      "version": "1.3.0-pre.2",
       "depth": 1,
       "source": "registry",
       "dependencies": {
diff --git a/com.unity.ml-agents/CHANGELOG.md b/com.unity.ml-agents/CHANGELOG.md
index 1af8571c56..5087ccc99a 100755
--- a/com.unity.ml-agents/CHANGELOG.md
+++ b/com.unity.ml-agents/CHANGELOG.md
@@ -22,6 +22,7 @@ and this project adheres to
 
 ### Bug Fixes
 #### com.unity.ml-agents / com.unity.ml-agents.extensions (C#)
+- Fixed missing tensor Dispose in ModelRunner (#)
 
 #### ml-agents / ml-agents-envs
 - Bumped numpy version to >=1.21.2,<1.24.0 (#5997)
diff --git a/com.unity.ml-agents/Runtime/Inference/ModelRunner.cs b/com.unity.ml-agents/Runtime/Inference/ModelRunner.cs
index d7a9f10a80..aaa6f85c13 100644
--- a/com.unity.ml-agents/Runtime/Inference/ModelRunner.cs
+++ b/com.unity.ml-agents/Runtime/Inference/ModelRunner.cs
@@ -150,6 +150,10 @@ public void Dispose()
             if (m_Engine != null)
                 m_Engine.Dispose();
             m_TensorAllocator?.Reset(false);
+            foreach (var (name, tensor) in m_InputsByName)
+            {
+                tensor.Dispose();
+            }
         }
 
         void FetchSentisOutputs(string[] names)