GroupedData.apply pandas_udf crashing

from pyspark.sql import SparkSession
from pyspark.sql import functions as F


spark = (
    SparkSession
    .builder
    .appName("pandas_udf")
    .getOrCreate()
)

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],
    ("id", "v")
)

@F.pandas_udf("id long, v double", F.PandasUDFType.GROUPED_MAP)
def normalize(pdf):
    v = pdf.v
    return pdf.assign(v=(v - v.mean()) / v.std())

(
    df
    .groupby("id")
    .apply(normalize)
    .show()
)