How to initialize the weights in pytorch with my own function? (cs231n PyTorch notebook)

In the cs231n pytorch notebook, they define random_weight function to return a tensor object that initialized by Kaiming normalization, and zero_weight function to return a tensor contains zero bias.

The random_weight implementation:

def random_weight(shape):
    if len(shape) == 2:  # FC weight
        fan_in = shape[0]
    else:
        fan_in = np.prod(shape[1:]) 
    w = torch.randn(shape, device=device, dtype=dtype) * np.sqrt(2. / fan_in)
    w.requires_grad = True
    return w

The zero_weight implementation:

def zero_weight(shape):
    return torch.zeros(shape, device=device, dtype=dtype, requires_grad=True)

Now the goal is to define a three-layer ConvNet using the pytorch sequential api, and we have to use the two above functions to initialize all the weights in the ConvNet.

I found some implementation like this:

model = nn.Sequential(....)

def init_weights(layer):
    if type(layer) in [nn.Conv2d, nn.Linear]:
        random_weight(layer.weight.size())
        zero_weight(layer.bias.size())

model.apply(init_weights)
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, nesterov=True)

But I suspect that the weights are not initialized by the two functions. Because I print the layer.weight.data before and after calling the random_weight, I got same data.

So I try to modify the init_weights code like this:

def init_weights(layer):
    if type(layer) in [nn.Conv2d, nn.Linear]:
        layer.weight.data = random_weight(layer.weight.size())
        layer.bias.data = zero_weight(layer.bias.size())

But the performance really sucks, I get about 10% accuracy on the validation set. So, does anyone know the right way to using this two functions to initialize the weights?

All suggestions are welcome! Thanks a lot!