Hi ,
I am implementing torch.nn.functional.normalize from scratch , found source  code of it implementing ,
def normalize(input, p=2, dim=1, eps=1e-12, out=None):
    if not torch.jit.is_scripting():
        if type(input) is not Tensor and has_torch_function((input,)):
            return handle_torch_function(
                normalize, (input,), input, p=p, dim=dim, eps=eps, out=out)
    if out is None:
        denom = input.norm(p, dim, keepdim=True).clamp_min(eps).expand_as(input)
        return input / denom
    else:
        denom = input.norm(p, dim, keepdim=True).clamp_min_(eps).expand_as(input)
        return torch.div(input, denom, out=out)

I found it little bit hard, I have tried to built it as
def normalize(input):
    return input / np.sqrt(np.sum(input*input, axis=-1, keepdims=True))

do you think if it will work? so that i could just use as  same as normalize function in pytorch.
normalize(nn.Dense(...)(input))