Skip to content
Snippets Groups Projects
Commit 744ae813 authored by Jerome Hue's avatar Jerome Hue Committed by Olivier BICHLER
Browse files

chore: use forwardStep attribute in ProdConso model

parent 94230383
No related branches found
No related tags found
2 merge requests!279v0.4.0,!256Add a Stack operator
Pipeline #60138 passed
......@@ -18,7 +18,7 @@ class StackProdConso : public ProdConso {
Elts_t getRequiredMemory(
const IOIndex_t outputIdx,
const std::vector<DimSize_t> &inputsSize) const override final;
void resetConsummerProducer() override final;
void resetConsummerProducer() override;
};
class StackOpImpl : public OperatorImpl {
......
......@@ -32,11 +32,15 @@ Elts_t StackProdConso::getRequiredMemory(
const StackOp &op = dynamic_cast<const StackOp &>(mOp);
// The produced data after one forward pass is simply the input size,
// we do not produce the whole output tensor everytime.
return Elts_t::DataElts(op.getInput(inputIdx)->size());
if (op.forwardStep() <= op.maxElements()) {
return Elts_t::DataElts(op.getInput(inputIdx)->size());
} else {
return Elts_t::NoneElts();
}
}
void StackProdConso::resetConsummerProducer() {
ProdConso::updateConsummerProducer();
ProdConso::resetConsummerProducer();
const StackOp &op = dynamic_cast<const StackOp &>(mOp);
op.forwardStep() = 0;
......@@ -94,7 +98,6 @@ bool Aidge::StackOp::forwardDims(bool /*allowDataDependency*/) {
return false;
}
void StackOp::setBackend(const std::string &name, DeviceIdx_t device) {
if (Registrar<StackOp>::exists({name})) {
SET_IMPL_MACRO(StackOp, *this, name);
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment