diff --git a/src/optimizer/Optimizer.cpp b/src/optimizer/Optimizer.cpp index 723f1408cb6eedd1226e701777c2d8b00a1ff067..5e1e2f38a3401170ee2e39b6dd0f83ac12d65a5a 100644 --- a/src/optimizer/Optimizer.cpp +++ b/src/optimizer/Optimizer.cpp @@ -20,7 +20,9 @@ void Aidge::Optimizer::resetGrad(std::shared_ptr<GraphView> graphView) auto op = node->getOperator(); if (op->isAtomic()) { auto tensorOp = std::static_pointer_cast<OperatorTensor>(op); - tensorOp->getOutput(0)->grad()->zeros(); + for (auto outputTensor : tensorOp->getOutputs()) { + outputTensor->grad()->zeros(); + } } else { auto metaOp = std::static_pointer_cast<MetaOperator_Op>(op); resetGrad(metaOp->getMicroGraph());