diff --git a/runtime/lib/ttnn/runtime.cpp b/runtime/lib/ttnn/runtime.cpp index c0465f329..22cc537c1 100644 --- a/runtime/lib/ttnn/runtime.cpp +++ b/runtime/lib/ttnn/runtime.cpp @@ -510,23 +510,12 @@ Tensor getOpOutputTensor(OpContext opContextHandle, return createNullTensor(); } - ::ttnn::Tensor hostTensor = ::ttnn::from_device(*outPtr); - ::ttnn::Tensor outCopy = - ::ttnn::to_layout(hostTensor, ::ttnn::ROW_MAJOR_LAYOUT, std::nullopt, - std::nullopt, static_cast<::ttnn::IDevice *>(nullptr)); - - void *src = ::tt::tt_metal::get_raw_host_data_ptr(outCopy); - std::uint32_t outCopySize = outCopy.volume() * outCopy.element_size(); - std::shared_ptr data = ::tt::runtime::utils::malloc_shared(outCopySize); - std::memcpy(data.get(), src, outCopySize); - - auto tensor = std::make_shared<::ttnn::Tensor>( - ttnn::createStorage(data.get(), outCopy.volume(), - ::tt::target::DataType::Float32), - outCopy.shape().value, ::ttnn::DataType::FLOAT32, - ::ttnn::Layout::ROW_MAJOR); + std::shared_ptr<::ttnn::Tensor> hostTensor = + std::make_shared<::ttnn::Tensor>(::ttnn::to_layout( + ::ttnn::from_device(*outPtr), ::ttnn::Layout::ROW_MAJOR, std::nullopt, + std::nullopt, static_cast<::ttnn::IDevice *>(nullptr))); - return Tensor(std::static_pointer_cast(tensor), nullptr, + return Tensor(std::static_pointer_cast(hostTensor), nullptr, DeviceRuntime::TTNN); }