diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index f317dd7..e47aa4f 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -328,9 +328,11 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Shortcut *l) { //std::cout<<"convert Shortcut\n"; //std::cout<<"New plugin Shortcut\n"; - ITensor *tens = tensors[l->backLayer]; - IPlugin *plugin = new ShortcutRT(tens); - IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); + ITensor *back_tens = tensors[l->backLayer]; + IPlugin *plugin = new ShortcutRT(); + + ITensor *inputs[2] = { input, back_tens }; + IPluginLayer *lRT = networkRT->addPlugin(inputs, 2, *plugin); checkNULL(lRT); return lRT; } diff --git a/src/pluginsRT/ShortcutRT.cpp b/src/pluginsRT/ShortcutRT.cpp index faba490..25ebea8 100644 --- a/src/pluginsRT/ShortcutRT.cpp +++ b/src/pluginsRT/ShortcutRT.cpp @@ -4,8 +4,7 @@ class ShortcutRT : public IPlugin { public: - ShortcutRT(ITensor *tens) { - this->back_layer = tens; + ShortcutRT() { } ~ShortcutRT(){ @@ -41,9 +40,12 @@ public: virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override { dnnType *srcData = (dnnType*)reinterpret_cast(inputs[0]); + dnnType *srcDataBack = (dnnType*)reinterpret_cast(inputs[1]); dnnType *dstData = reinterpret_cast(outputs[0]); checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream)); + shortcutForward(srcDataBack, dstData, batchSize, c, h, w, 1, batchSize, c, h, w, 1); + return 0; } @@ -60,5 +62,4 @@ public: } int c, h, w; - ITensor *back_layer; };