From c49bd1865a6b0ad9efa2778ef90e61ba60cfc602 Mon Sep 17 00:00:00 2001 From: MohammadReza Hosseini Date: Sun, 21 Jun 2020 14:06:12 +0430 Subject: [PATCH] fixing reorg layer bug --- include/tkDNN/pluginsRT/ReorgRT.h | 6 +++--- src/Reorg.cpp | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/include/tkDNN/pluginsRT/ReorgRT.h b/include/tkDNN/pluginsRT/ReorgRT.h index ee85718..afd205d 100644 --- a/include/tkDNN/pluginsRT/ReorgRT.h +++ b/include/tkDNN/pluginsRT/ReorgRT.h @@ -21,9 +21,9 @@ public: } void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override { - c = inputDims[0].d[0]; - h = inputDims[0].d[1]; - w = inputDims[0].d[2]; + c = outputDims[0].d[0]; + h = outputDims[0].d[1]; + w = outputDims[0].d[2]; } int initialize() override { diff --git a/src/Reorg.cpp b/src/Reorg.cpp index 67ed011..fe37775 100644 --- a/src/Reorg.cpp +++ b/src/Reorg.cpp @@ -25,7 +25,7 @@ Reorg::~Reorg() { dnnType* Reorg::infer(dataDim_t &dim, dnnType* srcData) { - reorgForward(srcData, dstData, dim.n, dim.c, dim.h, dim.w, stride); + reorgForward(srcData, dstData, output_dim.n, output_dim.c, output_dim.h, output_dim.w, stride); dim = output_dim; return dstData;