Remove inconsistent tensor reshaping issue nv-morpheus#1955

dagardner-nv · Oct 17, 2024 · 485764c · 485764c
1 parent fa67db6
commit 485764c
Showing 1 changed file with 2 additions and 4 deletions.
diff --git 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