diff --git a/src/operator/Memorize.cpp b/src/operator/Memorize.cpp
index cd4a4808137a786b04d8d143e50b96b9648e7d9a..c4f0bc4bf7267d24264652d5ed6b0d50935e1aa4 100644
--- a/src/operator/Memorize.cpp
+++ b/src/operator/Memorize.cpp
@@ -76,7 +76,7 @@ void Aidge::Memorize_OpImpl::forward() {
 const std::string Aidge::Memorize_Op::Type = "Memorize";
 
 Aidge::Memorize_Op::Memorize_Op(const std::uint32_t endStep)
-    : OperatorTensor(Type, {InputCategory::Data, InputCategory::Param}, 2),
+    : OperatorTensor(Type, {InputCategory::Data, InputCategory::Data}, 2),
         mAttributes(std::make_shared<Attributes_>(
                     attr<MemorizeAttr::ScheduleStep>(0),
                     attr<MemorizeAttr::ForwardStep>(0),