# L2 normalization
    x = F.normalize(x, p=2, dim=1)

I think F.normalize cannot implement the original L2 norm as same as the paper, the original L2 norm should be dim=None, which is not supported by F.normalize