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 {

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;