You can define a function to initialize the weights based on the layer and just apply your model on it:
def weights_init(m):
    if isinstance(m, nn.Conv2d):
        torch.nn.init.xavier_uniform(m.weight.data, nn.init.calculate_gain('relu'))
        m.bias.data.zero_()
    elif isinstance(m, nn.Linear):
        ...

model.apply(weights_init)

You can find the different init functions here 108.