def loss_fn(x, y):
    x = F.normalize(x, dim=(- 1), p=2)
    y = F.normalize(y, dim=(- 1), p=2)
    return (2 - (2 * (x * y).sum(dim=(- 1))))
def loss_fn(x, y):
    x = F.normalize(x, dim=(- 1), p=2)
    y = F.normalize(y, dim=(- 1), p=2)
    return (2 - (2 * (x * y).sum(dim=(- 1))))