[CPU] support LayerNorm with 3D shape (#15075)
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
@@ -421,8 +421,9 @@ class LayerNorm(MultiPlatformOp):
|
||||
x: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
if _is_cpu_amx_available:
|
||||
bias_data = self.bias.data if self.use_bias else None
|
||||
return torch.ops.sgl_kernel.layernorm_cpu(
|
||||
x, self.weight.data, self.variance_epsilon
|
||||
x, self.weight.data, bias_data, self.variance_epsilon
|
||||
)
|
||||
else:
|
||||
return self.forward_native(x)
|
||||
|
||||
Reference in New Issue
Block a user