What is the use of detach(), I do not think it is necessary here. I think we can also use
emb = torch.nn.Embedding(4,2)
norm = emb.weight.norm(p=2, dim=1, keepdim=True)
emb.weight = emb.weight.div(norm.expand_as(emb.weight))

Is there any problems with the above snippet?