fix cudnnBatchNorm for 3D data

5500153a · chengduoZH · 4cb2966d · 5500153a
隐藏空白更改
内联并排

Showing with 2 addition and 2 deletion

paddle/gserver/layers/CudnnBatchNormLayer.cpp paddle/gserver/layers/CudnnBatchNormLayer.cpp +2 -2

未找到文件。
--- a/paddle/gserver/layers/CudnnBatchNormLayer.cpp
+++ b/paddle/gserver/layers/CudnnBatchNormLayer.cpp
@@ -37,7 +37,7 @@ bool CudnnBatchNormLayer::init(const LayerMap& layerMap,
 }

 void CudnnBatchNormLayer::reshape(int batchSize) {
-  hl_tensor_reshape(ioDesc_, batchSize, channels_, imageH_, imageW_);
+  hl_tensor_reshape(ioDesc_, batchSize, channels_, imageH_ * imageD_, imageW_);
 }

 void CudnnBatchNormLayer::forward(PassType passType) {
@@ -104,7 +104,7 @@ void CudnnBatchNormLayer::forward(PassType passType) {
                                   EPS,
                                   batchSize,
                                   channels_,
-                                   imageH_,
+                                   imageH_ * imageD_,
                                   imageW_);
    }
  }