diff --git a/src/operator/FC.cpp b/src/operator/FC.cpp index dd1d0577ead56a4c62b4e3e9ca567685221d0a12..26752a35947900874fac6b576c94a1879938d356 100644 --- a/src/operator/FC.cpp +++ b/src/operator/FC.cpp @@ -62,8 +62,13 @@ bool Aidge::FC_Op::forwardDims(bool /*allowDataDependency*/) { } // check optional bias if(getInput(2)) - AIDGE_ASSERT((getInput(2)->nbDims() == 1) && - (getInput(2)->template dims<1>()[0] == outChannels), + AIDGE_ASSERT((((getInput(2)->nbDims() == 1) && + (getInput(2)->template dims<1>()[0] == outChannels)) || + ((getInput(2)->nbDims() == 2)&& + (getInput(0)->nbDims() == 2)&& + (getInput(2)->template dims<2>()[0] == getInput(0)->template dims<2>()[0])&& + (getInput(2)->template dims<2>()[1] == outChannels) + )), "Wrong bias size for FC operator."); // <batch, OutChannels> mOutputs[0]->resize({getInput(0)->dims()[0], outChannels});