diff --git a/include/aidge/operator/Identity.hpp b/include/aidge/operator/Identity.hpp index 7348fa10a96c55914bae68983b5e3bd4a9c40b12..57cd20311a4e4c98966af0af98b9fe4533155ea6 100644 --- a/include/aidge/operator/Identity.hpp +++ b/include/aidge/operator/Identity.hpp @@ -40,7 +40,7 @@ public: static const std::string Type; Identity_Op() - : OperatorTensor(Type, 1, 0, 0) + : OperatorTensor(Type, 1, 0, 1) { mImpl = std::make_shared<OperatorImpl>(*this); } @@ -101,7 +101,10 @@ public: if (outputIdx >= nbInputs()) { AIDGE_THROW_OR_ABORT(std::runtime_error, "%s Operator has %hu outputs", type().c_str(), nbInputs()); } - return mInputs[outputIdx]; + if (mInputs[outputIdx] == nullptr){ + return mOutputs[outputIdx]; // Input is not initialized with empty tensor + } + return mInputs[outputIdx]; // Identity, so Output is Input } void setBackend(const std::string& /*name*/, DeviceIdx_t /*device*/ = 0) override final { // setBackend do nothing, Identity node has no backend it just pass the same Tensor