diff --git a/pyproject.toml b/pyproject.toml
index 30e5ab3..40c0a53 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "PyCXpress"
-version = "0.0.5"
+version = "0.0.6"
 description = "PyCXpress is a high-performance hybrid framework that seamlessly integrates Python and C++ to harness the flexibility of Python and the speed of C++ for efficient and expressive computation, particularly in the realm of deep learning and numerical computing."
 readme = "README.md"
 authors = ["chaoqing <chaoqingwang.nick@gmail.com>"]
diff --git a/src/PyCXpress/example/main.cpp b/src/PyCXpress/example/main.cpp
index 82a9fd3..d8c3255 100644
--- a/src/PyCXpress/example/main.cpp
+++ b/src/PyCXpress/example/main.cpp
@@ -8,23 +8,30 @@
 
 namespace pcx = PyCXpress;
 
-void show_test(pcx::Model &python) {
+void show_test(pcx::Model &model) {
     std::vector<double> data(12);
     for (size_t i = 0; i < 12; i++) {
         data[i] = i;
     }
 
-    std::vector<uint8_t> shape = {3, 4};
-    memcpy(python.set_buffer("data_to_be_reshaped", {12}), data.data(),
-           data.size() * sizeof(double));
-    memcpy(python.set_buffer("new_2d_shape", {2}), shape.data(),
+    std::vector<uint8_t> shape   = {3, 4};
+    void                *pBuffer = nullptr;
+    size_t               nBytes  = 0;
+    std::tie(pBuffer, nBytes) = model.set_buffer("data_to_be_reshaped", {12});
+    assert(data.size() * sizeof(double) == nBytes);
+
+    memcpy(pBuffer, data.data(), nBytes);
+    memcpy(model.set_buffer("new_2d_shape", {2}).first, shape.data(),
            shape.size() * sizeof(uint8_t));
 
-    python.run();
+    model.run();
 
+    // test retrieve output tensor
     void               *p = nullptr;
     std::vector<size_t> new_shape;
-    std::tie(p, new_shape) = python.get_buffer("output_a");
+    pcx::BufferPtr      buf;
+    std::tie(buf, new_shape) = model.get_buffer("output_a");
+    std::tie(p, nBytes)      = buf;
 
     std::cout << "output shape: ";
     std::copy(new_shape.begin(), new_shape.end(),
@@ -33,10 +40,17 @@ void show_test(pcx::Model &python) {
 
     size_t size = std::accumulate(new_shape.begin(), new_shape.end(), 1,
                                   std::multiplies<int>());
+    assert(nBytes == sizeof(double) * size);
     std::cout << "output data: ";
     std::copy((double *)p, (double *)p + size,
               std::ostream_iterator<double>(std::cout, ", "));
     std::cout << std::endl;
+
+    // test retrieve input tensor
+    std::tie(buf, new_shape) = model.get_buffer("new_2d_shape");
+    assert(new_shape.size() == 1 && new_shape.front() == shape.size());
+    assert(buf.second == shape.size() * sizeof(uint8_t));
+    assert(0 == std::memcmp(buf.first, shape.data(), buf.second));
 }
 
 int main(int argc, char *argv[]) {
diff --git a/src/PyCXpress/example/model.py b/src/PyCXpress/example/model.py
index 79b6c4e..4f8342a 100644
--- a/src/PyCXpress/example/model.py
+++ b/src/PyCXpress/example/model.py
@@ -91,7 +91,7 @@ def model(input: InputDataSet, output: OutputDataSet):
         with nullcontext():
             # print(input.data_to_be_reshaped)
             # print(input.new_2d_shape)
-            output.output_a = input.data_to_be_reshaped.reshape(input.new_2d_shape)
+            output.output_a = input.data_to_be_reshaped.reshape(input.new_2d_shape).T
             # print(output.output_a)
 
 
diff --git a/src/PyCXpress/include/PyCXpress/core.hpp b/src/PyCXpress/include/PyCXpress/core.hpp
index 531d5b3..e2272b0 100644
--- a/src/PyCXpress/include/PyCXpress/core.hpp
+++ b/src/PyCXpress/include/PyCXpress/core.hpp
@@ -24,101 +24,106 @@ namespace PyCXpress {
 namespace py = pybind11;
 using namespace utils;
 
-class PYCXPRESS_EXPORT Buffer {
+typedef std::pair<void *, size_t> BufferPtr;
+class PYCXPRESS_EXPORT            Buffer {
     typedef unsigned char Bytes;
 
     template <typename T>
     static py::array __to_array(const std::vector<size_t> &shape, void *data,
-                                size_t max_size) {
-        std::vector<size_t> stride(shape.size());
-        *stride.rbegin() = sizeof(T);
-        auto ps          = shape.rbegin();
-        for (auto pt = stride.rbegin() + 1; pt != stride.rend(); pt++, ps++) {
-            *pt = *(pt - 1) * (*ps);
+                                           size_t &max_size) {
+                   std::vector<size_t> stride(shape.size());
+                   *stride.rbegin() = sizeof(T);
+                   auto ps          = shape.rbegin();
+                   for (auto pt = stride.rbegin() + 1; pt != stride.rend(); pt++, ps++) {
+                       *pt = *(pt - 1) * (*ps);
         }
-        if (max_size < stride.front() * shape.front()) {
-            throw std::runtime_error("Buffer size is too small");
+                   auto real_size = stride.front() * shape.front();
+                   if (max_size < real_size) {
+                       throw std::runtime_error("Buffer size is too small");
         }
-        return py::array_t<T>{shape, std::move(stride), (T *)(data),
-                              py::none()};
+                   max_size = real_size;
+                   return py::array_t<T>{shape, std::move(stride), (T *)(data),
+                                         py::none()};
     }
 
 public:
     Buffer() : m_size(0), m_data(nullptr), m_converter(nullptr) {}
     Buffer(size_t size, const std::string &data_type) : m_size(size) {
-        m_data   = new Bytes[m_size];
-        m_length = size;
+                   m_data = new Bytes[m_size];
 
-        if (data_type == "bool") {
-            m_converter = __to_array<bool>;
-            m_length /= sizeof(bool);
+                   if (data_type == "bool") {
+                       m_converter = __to_array<bool>;
+                       m_itemsize  = sizeof(bool);
         } else if (data_type == "int8_t") {
-            m_converter = __to_array<int8_t>;
-            m_length /= sizeof(int8_t);
+                       m_converter = __to_array<int8_t>;
+                       m_itemsize  = sizeof(int8_t);
         } else if (data_type == "int16_t") {
-            m_converter = __to_array<int16_t>;
-            m_length /= sizeof(int16_t);
+                       m_converter = __to_array<int16_t>;
+                       m_itemsize  = sizeof(int16_t);
         } else if (data_type == "int32_t") {
-            m_converter = __to_array<int32_t>;
-            m_length /= sizeof(int32_t);
+                       m_converter = __to_array<int32_t>;
+                       m_itemsize  = sizeof(int32_t);
         } else if (data_type == "int64_t") {
-            m_converter = __to_array<int64_t>;
-            m_length /= sizeof(int64_t);
+                       m_converter = __to_array<int64_t>;
+                       m_itemsize  = sizeof(int64_t);
         } else if (data_type == "uint8_t") {
-            m_converter = __to_array<uint8_t>;
-            m_length /= sizeof(uint8_t);
+                       m_converter = __to_array<uint8_t>;
+                       m_itemsize  = sizeof(uint8_t);
         } else if (data_type == "uint16_t") {
-            m_converter = __to_array<uint16_t>;
-            m_length /= sizeof(uint16_t);
+                       m_converter = __to_array<uint16_t>;
+                       m_itemsize  = sizeof(uint16_t);
         } else if (data_type == "uint32_t") {
-            m_converter = __to_array<uint32_t>;
-            m_length /= sizeof(uint32_t);
+                       m_converter = __to_array<uint32_t>;
+                       m_itemsize  = sizeof(uint32_t);
         } else if (data_type == "uint64_t") {
-            m_converter = __to_array<uint64_t>;
-            m_length /= sizeof(uint64_t);
+                       m_converter = __to_array<uint64_t>;
+                       m_itemsize  = sizeof(uint64_t);
         } else if (data_type == "float") {
-            m_converter = __to_array<float>;
-            m_length /= sizeof(float);
+                       m_converter = __to_array<float>;
+                       m_itemsize  = sizeof(float);
         } else if (data_type == "double") {
-            m_converter = __to_array<double>;
-            m_length /= sizeof(double);
+                       m_converter = __to_array<double>;
+                       m_itemsize  = sizeof(double);
         } else if (data_type == "char") {
-            m_converter = __to_array<char>;
-            m_length /= sizeof(char);
+                       m_converter = __to_array<char>;
+                       m_itemsize  = sizeof(char);
         } else {
-            throw NotImplementedError(data_type);
+                       throw NotImplementedError(data_type);
         }
     }
     Buffer(Buffer &&ohs)
         : m_size(ohs.m_size),
-          m_length(ohs.m_length),
+          m_itemsize(ohs.m_itemsize),
           m_data(ohs.m_data),
           m_converter(ohs.m_converter) {
-        ohs.m_data = nullptr;
+                   ohs.m_data = nullptr;
     }
 
     ~Buffer() {
-        delete[] m_data;
-        m_data = nullptr;
+                   delete[] m_data;
+                   m_data = nullptr;
     }
 
-    void *set(const std::vector<size_t> &shape) {
-        m_array = m_converter(shape, m_data, m_size);
-        return m_data;
+    BufferPtr set(const std::vector<size_t> &shape) {
+                   auto real_size = m_size;
+                   m_array        = m_converter(shape, m_data, real_size);
+                   return std::make_pair(m_data, real_size);
     }
 
-    inline size_t itemsize() const { return m_size / m_length; }
+    inline size_t itemsize() const { return m_itemsize; }
 
     py::array &array() { return m_array; }
 
-    void reset() { m_array = m_converter({m_length}, m_data, m_size); }
+    void reset() {
+                   m_array = m_converter({m_size / m_itemsize}, m_data, m_size);
+    }
 
 private:
     size_t    m_size;
-    size_t    m_length;
+    size_t    m_itemsize;
     Bytes    *m_data;
     py::array m_array;
-    py::array (*m_converter)(const std::vector<size_t> &, void *, size_t);
+    py::array (*m_converter)(const std::vector<size_t> &, void *, size_t &);
 };
 
 class PYCXPRESS_EXPORT Model {
@@ -155,25 +160,35 @@ class PYCXPRESS_EXPORT Model {
     }
 
 
-    void *set_buffer(const std::string         &name,
-                     const std::vector<size_t> &shape) {
-        auto &buf = m_buffers[name];
-        void *p   = buf.set(shape);
+    BufferPtr set_buffer(const std::string         &name,
+                         const std::vector<size_t> &shape) {
+        auto &buf  = m_buffers[name];
+        auto  pBuf = buf.set(shape);
         m_input.attr("set_buffer_value")(name, buf.array());
-        return p;
+        return pBuf;
     }
 
-    std::pair<void *, std::vector<size_t>> get_buffer(const std::string &name) {
-        auto &array  = m_buffers[name].array();
-        auto  pShape = m_output_buffer_sizes.find(name);
-        if (pShape == m_output_buffer_sizes.end()) {
-            return std::make_pair(
-                array.request().ptr,
-                std::vector<size_t>(array.shape(),
-                                    array.shape() + array.ndim()));
-        } else {
-            return std::make_pair(array.request().ptr, pShape->second);
+    std::pair<BufferPtr, std::vector<size_t>> get_buffer(
+        const std::string &name) {
+        auto &buf   = m_buffers[name];
+        auto &array = buf.array();
+
+        auto                iter = m_output_buffer_sizes.find(name);
+        std::vector<size_t> shape;
+        if (iter != m_output_buffer_sizes.end()) {
+            shape = iter->second;
+        } else {  // must be an input tensor, we do not suggest retrieve the
+                  // data from this interface
+            auto p = array.shape();
+            shape.resize(array.ndim());
+            std::copy_n(p, shape.size(), shape.begin());
         }
+
+        auto nBytes = buf.itemsize() *
+                      std::accumulate(shape.begin(), shape.end(), size_t(1),
+                                      std::multiplies<size_t>());
+        return std::make_pair(std::make_pair(array.request().ptr, nBytes),
+                              std::move(shape));
     }
 
     void run() {