fix bug

microsoft · Oct 2, 2024 · f93ccb7 · f93ccb7
1 parent 2895f37
commit f93ccb7
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc b/onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc
@@ -173,7 +173,9 @@ Status LayerNormImpl::ComputeImpl(OpKernelContext* p_ctx, int64_t orig_axis, flo
   U* mean_data = nullptr;
   if (!simplified) {
     Tensor* mean = p_ctx->Output(output_index++, TensorShape(mean_inv_std_dev_dim));
-    mean_data = mean->MutableData<U>();
+    if (mean != nullptr) {
+      mean_data = mean->MutableData<U>();
+    }
   }
 
   U* inv_std_dev_data = nullptr;