g->accept(&shapeInference);
- auto input = loadInput(g->getInput(params["input"])->getOperation()->getOutputShape(0));
+ // Check nodes
+ auto inputNode = g->getInput(params["input"]);
+ if (inputNode == nullptr) {
+ throw PluginException("input node <" + params["input"] +"> not found" );
+ }
+
+ auto outputNode = g->getOutput(params["output"]);
+ if (outputNode == nullptr) {
+ throw PluginException("output node <" + params["output"] +"> not found" );
+ }
+
+ auto input = loadInput(inputNode->getOperation()->getOutputShape(0));
interpreter.setInput(params["input"], input);
g->accept(&interpreter);
- _out = new TensorVariant(interpreter.getResult(g->getOutput(params["output"]))[0]);
+ _out = new TensorVariant(interpreter.getResult(outputNode)[0]);
return _out;
}