如何将权重归一化添加到PyTorch的预训练VGG16中？

import torch
import torchvision.models as models
import torch.nn as nn

vgg16 = models.vgg16(pretrained=True)

for param in vgg16.parameters():
    param.requires_grad = False

num_features = vgg16.classifier[6].in_features
vgg16.classifier[6] = nn.Linear(num_features, num_classes)

def weight_normalize(model, p=2, dim=1):
    norm = torch.norm(model.weight.data, p, dim, keepdim=True)
    model.weight.data = torch.div(model.weight.data, norm)

weight_normalize(vgg16.classifier[6])