diff --git a/include/aidge/backend/cpu/data/GetCPUPtr.h b/include/aidge/backend/cpu/data/GetCPUPtr.h index 47e3b07e8fa08cdcd714745a9a49bb03e30f79f5..b3e0fd967457585e7e3719f92dde7d6d93eee903 100644 --- a/include/aidge/backend/cpu/data/GetCPUPtr.h +++ b/include/aidge/backend/cpu/data/GetCPUPtr.h @@ -12,13 +12,16 @@ #ifndef AIDGE_CPU_DATA_GETCPUPTR_H_ #define AIDGE_CPU_DATA_GETCPUPTR_H_ +#include <cstddef> +#include <memory> + #include "aidge/data/Tensor.hpp" namespace Aidge { -inline void *getCPUPtr(std::shared_ptr<Aidge::Data> const &data) { +inline void *getCPUPtr(std::shared_ptr<Aidge::Data> const &data, const std::size_t offset = 0) { const auto tensor = std::static_pointer_cast<Tensor>(data); - return tensor->getImpl()->hostPtr(tensor->getImplOffset()); + return tensor->getImpl()->hostPtr(tensor->getImplOffset() + offset); } } // namespace Aidge -#endif // AIDGE_CPU_DATA_GETCPUPTR_H_ \ No newline at end of file +#endif // AIDGE_CPU_DATA_GETCPUPTR_H_