Yes, it works for dim=None, in weight_norm, also, for default dim=0,
I used this formula,
lin.weight_g*(lin.weight_v/lin.weight_v.norm(dim=1, keepdim=True))

or even this works
lin.weight_g*(lin.weight_v/torch.norm_except_dim(lin.weight_v, 2, dim=0))

I found out that incase of gradient descent here, weight_g, weight_v get updated, and after optimizer.step(), weight stays as it is. (so maybe, last training batch will have no impact on weights?)
weight is calculated from these updated weight_g, weight_v, before forward.
also, the first lin.weight_g is calculated as,
torch.norm_except_dim(lin.weight, 2, dim=0)

which leads to an output, whose shape is not changed for the specified dimension.
also, the first lin.weight_v value is same as lin.weight value.
I think main thing here is,


torch.norm(dim=1) gives output, whose shape is changed only in the specified dimension, and if we keepdim=True, then shape of that dimension becomes 1.


torch.norm_except_dim(dim=1) gives output, whose shape is not changed for the specified dimension, and rest all dimensions shape becomes 1.


and we want lin.weight_g, lin.weight_v.norm shape to be same, therefore, we use, lin.weight_v.norm(dim=1, keepdim=True) or torch.norm_except_dim(lin.weight_v, 2, dim=0)


if I set dim=1 in weight_norm, then lin.weight_g would be calculated by torch.norm_except_dim(lin.weight, 2, dim=1), and lin.weight would consider lin.weight_v.norm(dim=0, keepdim=True) or torch.norm_except_dim(lin.weight_v, 2, dim=1) to match with shape of lin.weight_g