diff --git a/python/morpheus/morpheus/messages/memory/tensor_memory.py b/python/morpheus/morpheus/messages/memory/tensor_memory.py
index ce193cd5da..33c6d4f2bd 100644
--- a/python/morpheus/morpheus/messages/memory/tensor_memory.py
+++ b/python/morpheus/morpheus/messages/memory/tensor_memory.py
@@ -176,7 +176,5 @@ def set_tensor(self, name: str, tensor: NDArrayType):
         ValueError
             If the number of rows in `tensor` does not match `count`
         """
-        # Ensure that we have 2D array here (`ensure_2d` inserts the wrong axis)
-        reshaped_tensor = tensor if tensor.ndim == 2 else np.reshape(tensor, (tensor.shape[0], -1))
-        self._check_tensor(reshaped_tensor)
-        self._tensors[name] = reshaped_tensor
+        self._check_tensor(tensor)
+        self._tensors[name] = tensor