From 0155f62b443ad9b1bad7daa8d44ab130554a51c8 Mon Sep 17 00:00:00 2001
From: Pavel Glagolev
Date: Mon, 18 Aug 2025 15:27:17 +0300
Subject: [PATCH] Fix visual observation tensor indexing for Unity inference.
This change corrects the tensor indexing calculation in TensorExtensions.Index()
to properly support CHW (channels-height-width) format used by both Unity's
observation writers and ONNX models during inference.
---
com.unity.ml-agents/Runtime/Inference/TensorExtensions.cs | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/com.unity.ml-agents/Runtime/Inference/TensorExtensions.cs b/com.unity.ml-agents/Runtime/Inference/TensorExtensions.cs
index 9e358cf157..94a717eea5 100644
--- a/com.unity.ml-agents/Runtime/Inference/TensorExtensions.cs
+++ b/com.unity.ml-agents/Runtime/Inference/TensorExtensions.cs
@@ -57,9 +57,9 @@ public static int Index(this TensorShape shape, int n, int c, int h, int w)
{
int index =
n * shape.Height() * shape.Width() * shape.Channels() +
- h * shape.Width() * shape.Channels() +
- w * shape.Channels() +
- c;
+ c * shape.Height() * shape.Width() +
+ h * shape.Width() +
+ w;
return index;
}
}