diff --git a/src/operator/MetaOperator.cpp b/src/operator/MetaOperator.cpp index d93d7d320f0f508b20714943ae3c8ed7fc561ec8..060c725482fedf4d6093e5acb988b2c721c27edc 100644 --- a/src/operator/MetaOperator.cpp +++ b/src/operator/MetaOperator.cpp @@ -87,12 +87,26 @@ void Aidge::MetaOperator_Op::setBackend(const std::string &name, Aidge::DeviceId if (Registrar<MetaOperator_Op>::exists({name, type()})) { // A custom implementation exists for this meta operator mImpl = Registrar<MetaOperator_Op>::create({name, type()})(*this); + + // Set backend for in/out tensor of the MetaOp + for(auto i: mGraph->inputNodes()){ + auto op_i = std::static_pointer_cast<OperatorTensor>(i->getOperator()); + for(std::size_t in_idx=0; in_idx < op_i->nbInputs(); ++in_idx){ + op_i->getInput(in_idx)->setBackend(name, device); + } + } + for(auto o: mGraph->outputNodes()){ + auto op_o = std::static_pointer_cast<OperatorTensor>(o->getOperator()); + for(std::size_t out_idx=0; out_idx < op_o->nbOutputs(); ++out_idx){ + op_o->getOutput(out_idx)->setBackend(name, device); + } + } + }else{ + // Input/output tensors backend are also updated here. + mGraph->setBackend(name, device); } - // The micro-graph should always be set to the right backend, since it - // shares input/output tensors. - // Input/output tensors backend are updated here. - mGraph->setBackend(name, device); + } std::set<std::string> Aidge::MetaOperator_Op::getAvailableBackends() const {