File tree Expand file tree Collapse file tree 1 file changed +7
-2
lines changed
paddle/fluid/inference/tensorrt/convert Expand file tree Collapse file tree 1 file changed +7
-2
lines changed Original file line number Diff line number Diff line change @@ -23,7 +23,7 @@ class ReluOpConverter : public OpConverter {
23
23
public:
24
24
ReluOpConverter () {}
25
25
void operator ()(const framework::proto::OpDesc& op,
26
- const framework::Scope& scope) override {
26
+ const framework::Scope& scope, bool test_mode ) override {
27
27
// Here the two nullptr looks strange, that's because the
28
28
// framework::OpDesc's constructor is strange.
29
29
framework::OpDesc op_desc (op, nullptr );
@@ -34,7 +34,12 @@ class ReluOpConverter : public OpConverter {
34
34
nvinfer1::IActivationLayer* layer = TRT_ENGINE_ADD_LAYER (
35
35
engine_, Activation, *const_cast <nvinfer1::ITensor*>(input_tensor),
36
36
nvinfer1::ActivationType::kRELU );
37
- engine_->DeclareOutput (layer, 0 , op_desc.Output (" Out" )[0 ]);
37
+ auto output_name = op_desc.Output (" Out" )[0 ];
38
+ engine_->SetITensor (output_name, layer->getOutput (0 ));
39
+ if (test_mode) { // the test framework can not determine which is the
40
+ // output, so place the declaration inside.
41
+ engine_->DeclareOutput (output_name);
42
+ }
38
43
}
39
44
};
40
45
You can’t perform that action at this time.
0 commit comments