You could write a weight_init method and apply it on the model:
def weight_init(m):
    if isinstance(m, nn.Conv2d):
        print('initializing conv2d weight')
        torch.nn.init.xavier_uniform_(m.weight)

model.apply(weight_init)