[BUG]: Python impl of TensorMemory.set_tensor reshapes incoming tensors

[dagardner-nv]

Version

24.10

Which installation method(s) does this occur on?

No response

Describe the bug.

The Python impl of TensorMemory.set_tensor will reshape the incoming tensor to 2 dimensions.

The problem is the C++ impl doesn't perform this, neither does the Python impl of the TensorMemory constructor or TensorMemory.set_tensors.

Minimum reproducible example

import cupy as cp

from morpheus.config import CppConfig
from morpheus.messages.memory.tensor_memory import TensorMemory

def repo(use_cpp:bool):
    CppConfig.set_should_use_cpp(use_cpp)
    mem = TensorMemory(count=3, tensors={"input_ids": cp.array([1, 2, 3]), "input_mask": cp.array([1, 1, 1]), "segment_ids": cp.array([0, 0, 1])})
    mem.set_tensor('new_tensor', cp.array([4, 5, 6]))
    print(f"use_cpp: {use_cpp}\n{mem.get_tensors()}\n")


repo(True)
repo(False)

Relevant log output

Click here to see error details

use_cpp: True

{'input_ids': array([1, 2, 3]), 'input_mask': array([1, 1, 1]), 'new_tensor': array([4, 5, 6]), 'segment_ids': array([0, 0, 1])}
use_cpp: False

{'input_ids': array([1, 2, 3]), 'input_mask': array([1, 1, 1]), 'segment_ids': array([0, 0, 1]), 'new_tensor': array([[4],

[5],

[6]])}

Full env printout

Click here to see environment details

[Paste the results of print_env.sh here, it will be hidden by default]

Other/Misc.

No response

Code of Conduct

• I agree to follow Morpheus' Code of Conduct
• I have searched the open bugs and have found no duplicates for this bug report