diff --git a/paddle/fluid/operators/mkldnn/gaussian_random_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/gaussian_random_mkldnn_op.cc index 503d3ec33762fb45f4132d133c3c5ad5540eca1a..de999035fa5d877996f1cc73b13f8193e6b87687 100644 --- a/paddle/fluid/operators/mkldnn/gaussian_random_mkldnn_op.cc +++ b/paddle/fluid/operators/mkldnn/gaussian_random_mkldnn_op.cc @@ -16,6 +16,7 @@ limitations under the License. */ #include "paddle/fluid/framework/generator.h" #include "paddle/fluid/operators/fill_constant_op.h" +#include "paddle/fluid/platform/mkldnn_helper.h" namespace paddle { namespace operators { @@ -42,7 +43,7 @@ class GaussianMKLDNNKernel : public paddle::framework::OpKernel { } tensor->set_layout(DataLayout::kMKLDNN); - tensor->set_format(dnnl::memory::format_tag::oihw); + tensor->set_format(platform::GetPlainMKLDNNFormat(tensor->dims().size())); } }; } // namespace operators