From 485764c7ceaf2de1352ec16ae3495ecc0fb894b1 Mon Sep 17 00:00:00 2001 From: David Gardner Date: Thu, 17 Oct 2024 16:18:27 -0700 Subject: [PATCH] Remove inconsistent tensor reshaping issue #1955 --- python/morpheus/morpheus/messages/memory/tensor_memory.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) 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