This should work:
batch_size, seq_size, dim = 2, 3, 4
embedding = torch.randn(batch_size, seq_size, dim)
print("x: ", embedding)

layer_norm = torch.nn.LayerNorm(dim)
y = layer_norm(embedding)
print("y: ", y)

out = (embedding - torch.mean(embedding, dim=2, keepdims=True)) / torch.sqrt(torch.var(embedding, dim=2, keepdims=True, unbiased=False) + layer_norm.eps)

print((out - y).abs().max())
# > tensor(1.1921e-07, grad_fn=<MaxBackward1>)